DP/Value Iteration Solution.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pprint\n",
    "import sys\n",
    "if \"../\" not in sys.path:\n",
    "  sys.path.append(\"../\") \n",
    "from lib.envs.gridworld import GridworldEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp = pprint.PrettyPrinter(indent=2)\n",
    "env = GridworldEnv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def value_iteration(env, theta=0.0001, discount_factor=1.0):\n",
    "    \"\"\"\n",
    "    Value Iteration Algorithm.\n",
    "    \n",
    "    Args:\n",
    "        env: OpenAI env. env.P represents the transition probabilities of the environment.\n",
    "            env.P[s][a] is a list of transition tuples (prob, next_state, reward, done).\n",
    "            env.nS is a number of states in the environment. \n",
    "            env.nA is a number of actions in the environment.\n",
    "        theta: We stop evaluation once our value function change is less than theta for all states.\n",
    "        discount_factor: Gamma discount factor.\n",
    "        \n",
    "    Returns:\n",
    "        A tuple (policy, V) of the optimal policy and the optimal value function.\n",
    "    \"\"\"\n",
    "    \n",
    "    def one_step_lookahead(state, V):\n",
    "        \"\"\"\n",
    "        Helper function to calculate the value for all action in a given state.\n",
    "        \n",
    "        Args:\n",
    "            state: The state to consider (int)\n",
    "            V: The value to use as an estimator, Vector of length env.nS\n",
    "        \n",
    "        Returns:\n",
    "            A vector of length env.nA containing the expected value of each action.\n",
    "        \"\"\"\n",
    "        A = np.zeros(env.nA)\n",
    "        for a in range(env.nA):\n",
    "            for prob, next_state, reward, done in env.P[state][a]:\n",
    "                A[a] += prob * (reward + discount_factor * V[next_state])\n",
    "        return A\n",
    "    \n",
    "    V = np.zeros(env.nS)\n",
    "    while True:\n",
    "        # Stopping condition\n",
    "        delta = 0\n",
    "        # Update each state...\n",
    "        for s in range(env.nS):\n",
    "            # Do a one-step lookahead to find the best action\n",
    "            A = one_step_lookahead(s, V)\n",
    "            best_action_value = np.max(A)\n",
    "            # Calculate delta across all states seen so far\n",
    "            delta = max(delta, np.abs(best_action_value - V[s]))\n",
    "            # Update the value function. Ref: Sutton book eq. 4.10. \n",
    "            V[s] = best_action_value        \n",
    "        # Check if we can stop \n",
    "        if delta < theta:\n",
    "            break\n",
    "    \n",
    "    # Create a deterministic policy using the optimal value function\n",
    "    policy = np.zeros([env.nS, env.nA])\n",
    "    for s in range(env.nS):\n",
    "        # One step lookahead to find the best action for this state\n",
    "        A = one_step_lookahead(s, V)\n",
    "        best_action = np.argmax(A)\n",
    "        # Always take the best action\n",
    "        policy[s, best_action] = 1.0\n",
    "    \n",
    "    return policy, V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Policy Probability Distribution:\n",
      "[[1. 0. 0. 0.]\n",
      " [0. 0. 0. 1.]\n",
      " [0. 0. 0. 1.]\n",
      " [0. 0. 1. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [0. 0. 1. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [0. 1. 0. 0.]\n",
      " [0. 0. 1. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [0. 1. 0. 0.]\n",
      " [0. 1. 0. 0.]\n",
      " [1. 0. 0. 0.]]\n",
      "\n",
      "Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):\n",
      "[[0 3 3 2]\n",
      " [0 0 0 2]\n",
      " [0 0 1 2]\n",
      " [0 1 1 0]]\n",
      "\n",
      "Value Function:\n",
      "[ 0. -1. -2. -3. -1. -2. -3. -2. -2. -3. -2. -1. -3. -2. -1.  0.]\n",
      "\n",
      "Reshaped Grid Value Function:\n",
      "[[ 0. -1. -2. -3.]\n",
      " [-1. -2. -3. -2.]\n",
      " [-2. -3. -2. -1.]\n",
      " [-3. -2. -1.  0.]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "policy, v = value_iteration(env)\n",
    "\n",
    "print(\"Policy Probability Distribution:\")\n",
    "print(policy)\n",
    "print(\"\")\n",
    "\n",
    "print(\"Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):\")\n",
    "print(np.reshape(np.argmax(policy, axis=1), env.shape))\n",
    "print(\"\")\n",
    "\n",
    "print(\"Value Function:\")\n",
    "print(v)\n",
    "print(\"\")\n",
    "\n",
    "print(\"Reshaped Grid Value Function:\")\n",
    "print(v.reshape(env.shape))\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the value function\n",
    "expected_v = np.array([ 0, -1, -2, -3, -1, -2, -3, -2, -2, -3, -2, -1, -3, -2, -1,  0])\n",
    "np.testing.assert_array_almost_equal(v, expected_v, decimal=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}