Control Algorithms 11.04.2021.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% Imports\n",
    "import gym\n",
    "import numpy as np\n",
    "# np.random.seed(123)\n",
    "import tensorflow as tf\n",
    "import time\n",
    "\n",
    "import sys\n",
    "# sys.path.append('../../')\n",
    "import algorithmsv2\n",
    "import cartpole_reward\n",
    "import estimate_L\n",
    "import observables\n",
    "import tf_algorithmsv2\n",
    "\n",
    "#%% Functions\n",
    "# Uniform distribution for importance sampling\n",
    "def rho(u, o='unif', a=0, b=1):\n",
    "    if o == 'unif':\n",
    "        return 1 / ( b - a )\n",
    "    if o == 'normal':\n",
    "        return np.exp( -u**2 / 2 ) / ( np.sqrt( 2 * np.pi ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading datasets\n",
    "X_0 = np.load('./random-agent/cartpole-states-0.npy').T\n",
    "X_1 = np.load('./random-agent/cartpole-states-1.npy').T\n",
    "Y_0 = np.load('./random-agent/cartpole-next-states-0.npy').T\n",
    "Y_1 = np.load('./random-agent/cartpole-next-states-1.npy').T\n",
    "\n",
    "# Creating Dictionaries for each dataset\n",
    "X_data = { 0: X_0, 1: X_1 }\n",
    "Y_data = { 0: Y_0, 1: Y_1 }\n",
    "\n",
    "X = np.append(X_data[0], X_data[1], axis=1)\n",
    "Y = np.append(Y_data[0], Y_data[1], axis=1)\n",
    "U = np.empty([1,X.shape[1]])\n",
    "for i in range(X_data[0].shape[1]):\n",
    "    U[:,i] = [0]\n",
    "for i in range(X_data[1].shape[1]):\n",
    "    U[:,i+X_data[0].shape[1]] = [1]\n",
    "\n",
    "dim_x = X.shape[0] # dimension of each data point (snapshot)\n",
    "dim_u = U.shape[0] # dimension of each action\n",
    "N = X.shape[1] # number of data points (snapshots)\n",
    "\n",
    "#%% Matrix builder functions\n",
    "order = 2\n",
    "phi = observables.monomials(order)\n",
    "psi = observables.monomials(order)\n",
    "\n",
    "#%% Compute Phi and Psi matrices + dimensions\n",
    "Phi_X = phi(X)\n",
    "Phi_Y = phi(Y)\n",
    "Psi_U = psi(U)\n",
    "\n",
    "dim_phi = Phi_X.shape[0]\n",
    "dim_psi = Psi_U.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us seek a finite dimensional approximation of the Koopman operator $\\mathcal{K}^{u_i}$.  Denote $K \\in \\mathbb{R}^{d\\times d \\times d}$ as a 3-d tensor. For any $u$, let us denote $K^{u}\\in\\mathbb{R}^{d\\times d}$ as follows: $K^{u}[i,j] = \\sum_{z = 1}^d K(i, j, z) \\psi(u)[z]$. Namely, $K^u$ is the result of the tensor vector product along the 3-d dimension of $K$ and $K^u$ serves as the finite dimensional approximation of Koopman operator $\\mathcal{K}^u$.\n",
    "We learn $K$ as follows:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\min_{K} \\sum_{i=1}^N \\left\\|  K^{u_i} \\phi(x_i) - \\phi(x_i')  \\right\\|^2.\n",
    "\\end{align}\n",
    "$$\n",
    "We can slightly re-write the above objective so that it becomes the regular multi-variate linear regression problem.  We can rearrange to write $K$ as a 2-d dimension matrix in $\\mathbb{R}^{d\\times d^2}$. Denote $M \\in \\mathbb{R}^{d\\times d^2}$, where $M[i, :] \\in \\mathbb{R}^{d^2}$ is the vector from stacking the columns of the 2-d matrix $K[i, :, :]$. Denote $ \\psi(u)\\otimes \\phi(x)\\in\\mathbb{R}^{d^2}$ as the Kronecker product. Then by linear algebra, we have:\n",
    "$\n",
    "\\begin{align*}\n",
    "K^{u} \\phi(x) = M ( \\psi(u)\\otimes \\phi(x) ).\n",
    "\\end{align*}$\n",
    "Thus the optimization problem becomes a regular linear regression:\n",
    "$\n",
    "\\begin{align*}\n",
    "\\min_{M} \\sum_{i=1}^N \\left\\|  M \\left( \\psi(u_i)\\otimes  \\phi(x_i)\\right) - \\phi(x_i')  \\right\\|^2,\n",
    "\\end{align*}$\n",
    " i.e., we do regression from vector $\\psi(u)\\otimes  \\phi(x_i) \\in \\mathbb{R}^{d^2}$ to $\\phi(x_i') \\in \\mathbb{R}^{d}$.\n",
    "\n",
    "To ensure data across different dimension share information, we should actually do a rank constraint here, i.e., we perform reduced-rank regression:\n",
    "$\n",
    "\\begin{align*}\n",
    "& \\min_{M} \\sum_{i=1}^N \\left\\|  M \\left( \\psi(u_i)\\otimes  \\phi(x_i)\\right) - \\phi(x_i')  \\right\\|^2,\\\\\n",
    "& \\text{s.t., } \\text{rank}(M) \\leq r,\n",
    "\\end{align*}$\n",
    " where $r < d$ is some hyper-parameter.\n",
    "\n",
    "Compute estimate of K tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "M shape: (15, 45)\n"
     ]
    }
   ],
   "source": [
    "#%% Build kronMatrix\n",
    "kronMatrix = np.empty((dim_psi * dim_phi, N))\n",
    "for i in range(N):\n",
    "    kronMatrix[:,i] = np.kron(Psi_U[:,i], Phi_X[:,i])\n",
    "\n",
    "#%% Estimate M and B matrices\n",
    "M = estimate_L.ols(kronMatrix.T, Phi_Y.T).T\n",
    "print(\"M shape:\", M.shape)\n",
    "assert M.shape == (dim_phi, dim_phi * dim_psi)\n",
    "\n",
    "B = estimate_L.ols(Phi_X.T, X.T)\n",
    "assert B.shape == (dim_phi, X.shape[0])\n",
    "\n",
    "#%% Reshape M into K tensor\n",
    "K = np.empty((dim_phi, dim_phi, dim_psi))\n",
    "for i in range(dim_phi):\n",
    "    K[i] = M[i].reshape((dim_phi,dim_psi), order='F')\n",
    "\n",
    "def K_u(K, u):\n",
    "    if len(u.shape) == 1:\n",
    "        u = u.reshape(-1,1) # assume transposing row vector into column vector\n",
    "    # u must be column vector\n",
    "    return np.einsum('ijz,z->ij', K, psi(u)[:,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Control Algorithms 10/28/2021\n",
    "### Notebook that covers our current implementations of control and aims to be used in fixing the problems with the gradient explosion\n",
    "\n",
    "## Learning Optimal Policy via Directly Minimizing Bellman Error\n",
    "This is based on section 5 in the [overleaf writeup](https://www.overleaf.com/project/6155ef2f57f2b6a1e034b696)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "w: [[1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [1.]]\n",
      "tf_w: <tf.Variable 'weights:0' shape=(15,) dtype=float32, numpy=\n",
      "array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
      "      dtype=float32)>\n"
     ]
    }
   ],
   "source": [
    "All_U = np.array([[0,1]])\n",
    "u_bounds = [0,1]\n",
    "learning_rate = 1e-2\n",
    "w = np.ones([dim_phi,1])\n",
    "print(\"w:\", w)\n",
    "tf_w = tf.Variable(tf.ones([dim_phi]), trainable=True, name=\"weights\")\n",
    "print(\"tf_w:\", tf_w)\n",
    "batch_size = 256\n",
    "\n",
    "def cost(x,u):\n",
    "    return -cartpole_reward.defaultCartpoleReward(x,u)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Expressing Bellman Error:\n",
    "$$\n",
    "\\begin{align*}\n",
    "    (w^{\\top}\\phi(x) - \\min_{\\pi: x\\mapsto \\Delta(A)} \\left[  \\mathbb{E}_{u\\sim \\pi(x)} \\left[ c(x,u) + \\ln\\pi(u | x) + w^{\\top} K(I, I, \\psi(u)) \\phi(x) \\right] \\right])^2.\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "### Policy Expression:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\pi(u | x) = \\exp\\left( - \\left( c(x,u) + w^{\\top} K(I,I, \\psi(u)) \\phi(x)   \\right)  \\right) / Z_x,\n",
    "\\end{align}\n",
    "$$\n",
    "In the next block, we show our coded implementation of the policy expression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inner_pi_u(u, x):\n",
    "    inner_pi_u = (-(cost(x, u) + w.T @ K_u(K, u) @ phi(x)))[0]\n",
    "    return inner_pi_u\n",
    "\n",
    "def pi_u(u, x):\n",
    "    inner = inner_pi_u(u, x)\n",
    "    return np.exp(inner)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorFlow version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pi(u, x):\n",
    "    phi_x = tf.cast(tf.stack(phi(x)), tf.float32)\n",
    "    \n",
    "    @tf.autograph.experimental.do_not_convert\n",
    "    def compute_numerator(i):\n",
    "        u = All_U[:,i]\n",
    "        u = tf.reshape(u, [tf.shape(u)[0],1])\n",
    "\n",
    "        phi_x_prime = tf.tensordot(tf.cast(K_u(K, u), tf.float32), phi_x, axes=1)\n",
    "        weighted_phi_x_prime = tf.tensordot(tf_w, phi_x_prime, axes=1)\n",
    "        inner = tf.add(cost(x,u), weighted_phi_x_prime)\n",
    "        return tf.math.exp(-inner)\n",
    "    Z_x = tf.math.reduce_sum(tf.map_fn(fn=compute_numerator, elems=tf.range(All_U.shape[1]), dtype=tf.float32))\n",
    "\n",
    "    phi_x_prime = tf.tensordot(tf.cast(K_u(K, u), tf.float32), phi_x, axes=1)\n",
    "    weighted_phi_x_prime = tf.tensordot(tf_w, phi_x_prime, axes=1)\n",
    "    inner = tf.add(cost(x,u), weighted_phi_x_prime)\n",
    "    numerator = tf.math.exp(-inner)\n",
    "\n",
    "    pi_value = tf.divide(numerator, Z_x)\n",
    "\n",
    "    return pi_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Expressing Bellman Error over a dataset:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\min_{w: \\|w\\|_2 \\leq W} \\sum_{i=1}^N \\left( w^{\\top}\\phi(x) - \\min_{\\pi(\\cdot | s)} \\left[ \\mathbb{E}_{u\\sim \\pi(x)} \\left[ c(x,u) + \\ln\\pi(u|s) +  w^{\\top} K(I, I, \\psi(u)) \\phi(x) \\right] \\right] \\right)^2\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discreteBellmanError():\n",
    "    ''' Equation 12 in writeup '''\n",
    "    total = 0\n",
    "    for i in range(X.shape[1]):\n",
    "        x = X[:,i].reshape(-1,1)\n",
    "        phi_x = phi(x)\n",
    "\n",
    "        inner_pi_us = []\n",
    "        for u in All_U.T:\n",
    "            u = u.reshape(-1,1)\n",
    "            inner_pi_us.append(inner_pi_u(u, x))\n",
    "        inner_pi_us = np.real(inner_pi_us)\n",
    "        max_inner_pi_u = np.max(inner_pi_us)\n",
    "        pi_us = np.exp(inner_pi_us - max_inner_pi_u)\n",
    "        Z_x = np.sum(pi_us)\n",
    "\n",
    "        expectation_u = 0\n",
    "        pi_sum = 0\n",
    "        for i,u in enumerate(All_U.T):\n",
    "            u = u.reshape(-1,1)\n",
    "            pi = pi_us[i] / Z_x\n",
    "            assert pi >= 0\n",
    "            pi_sum += pi\n",
    "            expectation_u += (cost(x, u) + np.log(pi) + w.T @ K_u(K, u) @ phi_x) * pi\n",
    "        total += np.power((w.T @ phi_x - expectation_u), 2)\n",
    "        assert np.isclose(pi_sum, 1, rtol=1e-3, atol=1e-4)\n",
    "    return total"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorFlow version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discreteBellmanErrorTF(x):\n",
    "    @tf.autograph.experimental.do_not_convert\n",
    "    def computeError(i):\n",
    "        phi_x = tf.cast(tf.stack(phi(x)), tf.float32)\n",
    "\n",
    "        # Sample u1 and u2, get psi_u1 and psi_u2\n",
    "        random_index = tf.random.shuffle(tf.range(All_U.shape[1]))[0]\n",
    "        u1 = All_U[:,random_index]\n",
    "        u1 = tf.reshape(u1, [tf.shape(u1)[0],1])\n",
    "        random_index = tf.random.shuffle(tf.range(All_U.shape[1]))[0]\n",
    "        u2 = All_U[:,random_index]\n",
    "        u2 = tf.reshape(u2, [tf.shape(u2)[0],1])\n",
    "\n",
    "        # First term of value fn expressed in terms of dictionary\n",
    "        inner_part_1 = tf.tensordot(tf_w, phi_x, axes=1)\n",
    "        # Computing terms in RHS of Bellman eqn\n",
    "        cost_plus_log_pi = tf.cast(\n",
    "            tf.add(cost(x, u1), tf.math.log(pi(u1, x))),\n",
    "            tf.float32\n",
    "        )\n",
    "        phi_x_prime = tf.tensordot(tf.cast(K_u(K, u2), tf.float32), phi_x, axes=1)\n",
    "        weighted_phi_x_prime = tf.tensordot(tf_w, phi_x_prime, axes=1)\n",
    "\n",
    "        inner_part_2 = tf.add(cost_plus_log_pi, weighted_phi_x_prime)\n",
    "        importanceWeight = tf.cast(\n",
    "            tf.multiply(pi(u1, x), All_U.shape[1]),\n",
    "            tf.float32\n",
    "        )\n",
    "        inner_part_2 = tf.multiply(importanceWeight, inner_part_2)\n",
    "        \n",
    "        inner_difference = tf.subtract(inner_part_1, inner_part_2)\n",
    "        squared_inner = tf.math.square(inner_difference)\n",
    "        return squared_inner\n",
    "\n",
    "    results = tf.map_fn(fn=computeError, elems=tf.range(X.shape[1]/100), dtype=tf.float32)\n",
    "\n",
    "    return tf.math.reduce_sum(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to minimize the ojective function above, we will run some sort of gradient descent algorithm, in this case, SGD (stochastic gradient descent).\n",
    "\n",
    "The gradient of the objective function is as follows:\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\nabla_{w} := \\left( w^{\\top}\\phi(x) - \\left[ \\mathbb{E}_{u\\sim \\pi(x)} \\left[ c(x,u) + \\ln\\pi(u|s) +  w^{\\top} K(I, I, \\psi(u)) \\phi(x) \\right] \\right] \\right)\\left[ \\phi(x) -  \\mathbb{E}_{u\\sim \\pi(\\cdot | x)}  K(I, I, \\psi(u)) \\phi(x)  \\right].\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Our programatic implementation takes advantage of importance weighting as well as control sampling to get an unbiased estimate of the above formulation\n",
    "$$\n",
    "\\begin{align}\n",
    "\\widetilde\\nabla_{w} = &  \\left( w^{\\top}\\phi(x) - \\left[  \\frac{ \\pi(u_1|x) }{\\rho(u_1)} \\left[ c(x,u_1) + \\ln\\pi(u_1|s) +  w^{\\top} K(I, I, \\psi(u_1)) \\phi(x) \\right] \\right] \\right)\\\\\n",
    "& \\qquad  \\cdot \\left[ \\phi(x) -   \\frac{ \\pi(u_2|x) }{\\rho(u)}  K(I, I, \\psi(u_2)) \\phi(x)  \\right].\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Minibatch implementation\n",
    "x_batch = np.empty([X.shape[0],batch_size])\n",
    "for i in range(batch_size):\n",
    "    x = X[:, np.random.choice(np.arange(X.shape[1]))]\n",
    "    x_batch[:,i] = x\n",
    "phi_x_batch = phi(x_batch)\n",
    "\n",
    "nabla_w = np.zeros_like(w)\n",
    "for x1, phi_x1 in zip(x_batch.T, phi_x_batch.T):\n",
    "    x1 = x1.reshape(-1,1)\n",
    "    phi_x1 = phi_x1.reshape(-1,1)\n",
    "\n",
    "    inner_pi_us = []\n",
    "    for u in All_U.T:\n",
    "        u = u.reshape(-1,1)\n",
    "        inner_pi_us.append(inner_pi_u(u, x1))\n",
    "    inner_pi_us = np.real(inner_pi_us)\n",
    "    max_inner_pi_u = np.max(inner_pi_us)\n",
    "    pi_us = np.exp(inner_pi_us - max_inner_pi_u)\n",
    "    Z_x = np.sum(pi_us)\n",
    "\n",
    "    expectationTerm1 = 0\n",
    "    expectationTerm2 = 0\n",
    "    for i,u in enumerate(All_U.T):\n",
    "        u = u.reshape(-1,1)\n",
    "        K_u_const = K_u(K, u)\n",
    "        expectationTerm1 += (pi_us[i] / Z_x) * (cost(x1, u) + np.log(pi_us[i] / Z_x) + w.T @ K_u_const @ phi_x1)\n",
    "        expectationTerm2 += (pi_us[i] / Z_x) * K_u_const @ phi_x1\n",
    "\n",
    "    # Equation 13/14 in writeup\n",
    "    nabla_w += ((w.T @ phi_x1 - expectationTerm1) * (phi_x1 - expectationTerm2)) / batch_size\n",
    "\n",
    "# nabla_w = (w @ phi_x1 - ((cost(x1, u1) + np.log(pi_u(u1, x1)) + w @ K_u(K, u1) @ phi_x1))) \\\n",
    "#             * (phi_x1 - K_u(K, u2) @ phi_x1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update the weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.         0.99818444 1.00130054 1.0117048  0.96920474 1.00628324\n",
      " 0.98857944 0.99686526 1.07847905 0.91424755 0.99560384 0.85357811\n",
      " 1.10501536 0.86346989 0.59492226]\n",
      "21304.21022187679\n"
     ]
    }
   ],
   "source": [
    "w = w - (learning_rate * nabla_w)\n",
    "print(w[:,0])\n",
    "print(discreteBellmanError()[0,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorFlow version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-11-86a944c5e22d>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[1;31m# Compute loss\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      7\u001b[0m \u001b[1;32mwith\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mGradientTape\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtape\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 8\u001b[1;33m     \u001b[0mloss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mdiscreteBellmanErrorTF\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      9\u001b[0m \u001b[1;31m# Compute gradient\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     10\u001b[0m \u001b[0mgrads\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtape\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mgradient\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mtf_w\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-8-8424dfd53196>\u001b[0m in \u001b[0;36mdiscreteBellmanErrorTF\u001b[1;34m(x)\u001b[0m\n\u001b[0;32m     33\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0msquared_inner\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     34\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 35\u001b[1;33m     \u001b[0mresults\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmap_fn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mcomputeError\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0melems\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mX\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m/\u001b[0m\u001b[1;36m100\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     36\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     37\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmath\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mreduce_sum\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mresults\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\util\\deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    572\u001b[0m                   \u001b[0mfunc\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__module__\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0marg_name\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0marg_value\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'in a future version'\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    573\u001b[0m                   if date is None else ('after %s' % date), instructions)\n\u001b[1;32m--> 574\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    575\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    576\u001b[0m     doc = _add_deprecated_arg_value_notice_to_docstring(\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\util\\deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    505\u001b[0m                 \u001b[1;34m'in a future version'\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mdate\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m \u001b[1;32melse\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;34m'after %s'\u001b[0m \u001b[1;33m%\u001b[0m \u001b[0mdate\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    506\u001b[0m                 instructions)\n\u001b[1;32m--> 507\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    508\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    509\u001b[0m     doc = _add_deprecated_arg_notice_to_docstring(\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\map_fn.py\u001b[0m in \u001b[0;36mmap_fn_v2\u001b[1;34m(fn, elems, dtype, parallel_iterations, back_prop, swap_memory, infer_shape, name, fn_output_signature)\u001b[0m\n\u001b[0;32m    631\u001b[0m   \u001b[1;32mif\u001b[0m \u001b[0mfn_output_signature\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    632\u001b[0m     \u001b[0mfn_output_signature\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mdtype\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 633\u001b[1;33m   return map_fn(\n\u001b[0m\u001b[0;32m    634\u001b[0m       \u001b[0mfn\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mfn\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    635\u001b[0m       \u001b[0melems\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0melems\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\util\\deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    505\u001b[0m                 \u001b[1;34m'in a future version'\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mdate\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m \u001b[1;32melse\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;34m'after %s'\u001b[0m \u001b[1;33m%\u001b[0m \u001b[0mdate\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    506\u001b[0m                 instructions)\n\u001b[1;32m--> 507\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    508\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    509\u001b[0m     doc = _add_deprecated_arg_notice_to_docstring(\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\map_fn.py\u001b[0m in \u001b[0;36mmap_fn\u001b[1;34m(fn, elems, dtype, parallel_iterations, back_prop, swap_memory, infer_shape, name, fn_output_signature)\u001b[0m\n\u001b[0;32m    491\u001b[0m       \u001b[1;32mreturn\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m \u001b[1;33m+\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtas\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    492\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 493\u001b[1;33m     _, r_a = control_flow_ops.while_loop(\n\u001b[0m\u001b[0;32m    494\u001b[0m         \u001b[1;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mi\u001b[0m \u001b[1;33m<\u001b[0m \u001b[0mn\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    495\u001b[0m         \u001b[0mcompute\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult_batchable_ta\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\control_flow_ops.py\u001b[0m in \u001b[0;36mwhile_loop\u001b[1;34m(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)\u001b[0m\n\u001b[0;32m   2733\u001b[0m                                               list(loop_vars))\n\u001b[0;32m   2734\u001b[0m       \u001b[1;32mwhile\u001b[0m \u001b[0mcond\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mloop_vars\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2735\u001b[1;33m         \u001b[0mloop_vars\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mbody\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mloop_vars\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2736\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mtry_to_pack\u001b[0m \u001b[1;32mand\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mloop_vars\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_basetuple\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2737\u001b[0m           \u001b[0mpacked\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\control_flow_ops.py\u001b[0m in \u001b[0;36m<lambda>\u001b[1;34m(i, lv)\u001b[0m\n\u001b[0;32m   2724\u001b[0m         cond = lambda i, lv: (  # pylint: disable=g-long-lambda\n\u001b[0;32m   2725\u001b[0m             math_ops.logical_and(i < maximum_iterations, orig_cond(*lv)))\n\u001b[1;32m-> 2726\u001b[1;33m         \u001b[0mbody\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlv\u001b[0m\u001b[1;33m:\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m \u001b[1;33m+\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0morig_body\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mlv\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2727\u001b[0m       \u001b[0mtry_to_pack\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mFalse\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2728\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\map_fn.py\u001b[0m in \u001b[0;36mcompute\u001b[1;34m(i, tas)\u001b[0m\n\u001b[0;32m    481\u001b[0m       \u001b[0mag_ctx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mautograph_ctx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcontrol_status_ctx\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    482\u001b[0m       \u001b[0mautographed_fn\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mautograph\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtf_convert\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mag_ctx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 483\u001b[1;33m       \u001b[0mresult_value\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mautographed_fn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0melems_value\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    484\u001b[0m       \u001b[0mnest\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0massert_same_structure\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfn_output_signature\u001b[0m \u001b[1;32mor\u001b[0m \u001b[0melems\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult_value\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    485\u001b[0m       \u001b[0mresult_value_flat\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnest\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mflatten\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mresult_value\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\autograph\\impl\\api.py\u001b[0m in \u001b[0;36mwrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    300\u001b[0m   \u001b[1;32mdef\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    301\u001b[0m     \u001b[1;32mwith\u001b[0m \u001b[0mag_ctx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mControlStatusCtx\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mstatus\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mag_ctx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mStatus\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mDISABLED\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 302\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    303\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    304\u001b[0m   \u001b[1;32mif\u001b[0m \u001b[0minspect\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0misfunction\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mor\u001b[0m \u001b[0minspect\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mismethod\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-8-8424dfd53196>\u001b[0m in \u001b[0;36mcomputeError\u001b[1;34m(i)\u001b[0m\n\u001b[0;32m     16\u001b[0m         \u001b[1;31m# Computing terms in RHS of Bellman eqn\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     17\u001b[0m         cost_plus_log_pi = tf.cast(\n\u001b[1;32m---> 18\u001b[1;33m             \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0madd\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcost\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mu1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmath\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlog\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mpi\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mu1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     19\u001b[0m             \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     20\u001b[0m         )\n",
      "\u001b[1;32m<ipython-input-6-965e1c74caee>\u001b[0m in \u001b[0;36mpi\u001b[1;34m(u, x)\u001b[0m\n\u001b[0;32m     11\u001b[0m         \u001b[0minner\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0madd\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcost\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mu\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mweighted_phi_x_prime\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     12\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmath\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexp\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m-\u001b[0m\u001b[0minner\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 13\u001b[1;33m     \u001b[0mZ_x\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmath\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mreduce_sum\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmap_fn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mcompute_numerator\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0melems\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mAll_U\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     14\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     15\u001b[0m     \u001b[0mphi_x_prime\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtensordot\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcast\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mK_u\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mK\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mu\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mphi_x\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0maxes\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\util\\deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    572\u001b[0m                   \u001b[0mfunc\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__module__\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0marg_name\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0marg_value\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'in a future version'\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    573\u001b[0m                   if date is None else ('after %s' % date), instructions)\n\u001b[1;32m--> 574\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    575\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    576\u001b[0m     doc = _add_deprecated_arg_value_notice_to_docstring(\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\util\\deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    505\u001b[0m                 \u001b[1;34m'in a future version'\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mdate\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m \u001b[1;32melse\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;34m'after %s'\u001b[0m \u001b[1;33m%\u001b[0m \u001b[0mdate\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    506\u001b[0m                 instructions)\n\u001b[1;32m--> 507\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    508\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    509\u001b[0m     doc = _add_deprecated_arg_notice_to_docstring(\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\map_fn.py\u001b[0m in \u001b[0;36mmap_fn_v2\u001b[1;34m(fn, elems, dtype, parallel_iterations, back_prop, swap_memory, infer_shape, name, fn_output_signature)\u001b[0m\n\u001b[0;32m    631\u001b[0m   \u001b[1;32mif\u001b[0m \u001b[0mfn_output_signature\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    632\u001b[0m     \u001b[0mfn_output_signature\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mdtype\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 633\u001b[1;33m   return map_fn(\n\u001b[0m\u001b[0;32m    634\u001b[0m       \u001b[0mfn\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mfn\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    635\u001b[0m       \u001b[0melems\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0melems\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\util\\deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    505\u001b[0m                 \u001b[1;34m'in a future version'\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mdate\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m \u001b[1;32melse\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;34m'after %s'\u001b[0m \u001b[1;33m%\u001b[0m \u001b[0mdate\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    506\u001b[0m                 instructions)\n\u001b[1;32m--> 507\u001b[1;33m       \u001b[1;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    508\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    509\u001b[0m     doc = _add_deprecated_arg_notice_to_docstring(\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\map_fn.py\u001b[0m in \u001b[0;36mmap_fn\u001b[1;34m(fn, elems, dtype, parallel_iterations, back_prop, swap_memory, infer_shape, name, fn_output_signature)\u001b[0m\n\u001b[0;32m    491\u001b[0m       \u001b[1;32mreturn\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m \u001b[1;33m+\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtas\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    492\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 493\u001b[1;33m     _, r_a = control_flow_ops.while_loop(\n\u001b[0m\u001b[0;32m    494\u001b[0m         \u001b[1;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mi\u001b[0m \u001b[1;33m<\u001b[0m \u001b[0mn\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    495\u001b[0m         \u001b[0mcompute\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult_batchable_ta\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\control_flow_ops.py\u001b[0m in \u001b[0;36mwhile_loop\u001b[1;34m(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)\u001b[0m\n\u001b[0;32m   2732\u001b[0m       loop_var_structure = nest.map_structure(type_spec.type_spec_from_value,\n\u001b[0;32m   2733\u001b[0m                                               list(loop_vars))\n\u001b[1;32m-> 2734\u001b[1;33m       \u001b[1;32mwhile\u001b[0m \u001b[0mcond\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mloop_vars\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2735\u001b[0m         \u001b[0mloop_vars\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mbody\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mloop_vars\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2736\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mtry_to_pack\u001b[0m \u001b[1;32mand\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mloop_vars\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_basetuple\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\control_flow_ops.py\u001b[0m in \u001b[0;36m<lambda>\u001b[1;34m(i, lv)\u001b[0m\n\u001b[0;32m   2723\u001b[0m         \u001b[0mloop_vars\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mcounter\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mloop_vars\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2724\u001b[0m         cond = lambda i, lv: (  # pylint: disable=g-long-lambda\n\u001b[1;32m-> 2725\u001b[1;33m             math_ops.logical_and(i < maximum_iterations, orig_cond(*lv)))\n\u001b[0m\u001b[0;32m   2726\u001b[0m         \u001b[0mbody\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlv\u001b[0m\u001b[1;33m:\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m \u001b[1;33m+\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0morig_body\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mlv\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2727\u001b[0m       \u001b[0mtry_to_pack\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mFalse\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\map_fn.py\u001b[0m in \u001b[0;36m<lambda>\u001b[1;34m(i, _)\u001b[0m\n\u001b[0;32m    492\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    493\u001b[0m     _, r_a = control_flow_ops.while_loop(\n\u001b[1;32m--> 494\u001b[1;33m         \u001b[1;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mi\u001b[0m \u001b[1;33m<\u001b[0m \u001b[0mn\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    495\u001b[0m         \u001b[0mcompute\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult_batchable_ta\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    496\u001b[0m         \u001b[0mparallel_iterations\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mparallel_iterations\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mA:\\Users\\wrozw\\anaconda3\\envs\\conda-rl-venv\\lib\\site-packages\\tensorflow\\python\\ops\\gen_math_ops.py\u001b[0m in \u001b[0;36mless\u001b[1;34m(x, y, name)\u001b[0m\n\u001b[0;32m   4906\u001b[0m   \u001b[1;32mif\u001b[0m \u001b[0mtld\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mis_eager\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   4907\u001b[0m     \u001b[1;32mtry\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 4908\u001b[1;33m       _result = pywrap_tfe.TFE_Py_FastPathExecute(\n\u001b[0m\u001b[0;32m   4909\u001b[0m         \u001b[0m_ctx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_context_handle\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtld\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdevice_name\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m\"Less\"\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mname\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtld\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mop_callbacks\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   4910\u001b[0m         x, y)\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "optimizer = tf.keras.optimizers.SGD()\n",
    "\n",
    "# Sample x from X\n",
    "x1 = X[:, tf.random.shuffle(tf.range(X.shape[1]))[0]]\n",
    "x1 = tf.reshape(x1, [tf.shape(x1)[0],1])\n",
    "# Compute loss\n",
    "with tf.GradientTape() as tape:\n",
    "    loss = discreteBellmanErrorTF(x1)\n",
    "# Compute gradient\n",
    "grads = tape.gradient(loss, [tf_w])\n",
    "# Apply gradient (update weights)\n",
    "optimizer.apply_gradients(zip(grads, [tf_w]))\n",
    "print(tf_w)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0727fff862403d9dc13e66806b9c6776c4f51a1be1742a67d5c804a4cd3707be"
  },
  "kernelspec": {
   "display_name": "Python 3.8.8 64-bit ('conda-rl-venv': conda)",
   "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.8.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}