Frozen Lake PI base code.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function,\n",
    "                        unicode_literals)\n",
    "%matplotlib inline\n",
    "\n",
    "import time\n",
    "import seaborn\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "from gym.envs.toy_text.frozen_lake import LEFT, RIGHT, DOWN, UP\n",
    "from gym.envs.toy_text import frozen_lake, discrete\n",
    "\n",
    "import gym\n",
    "from gym.envs.registration import register\n",
    "\n",
    "register(\n",
    "    id='D4x4-FrozenLake-v0',\n",
    "    entry_point='gym.envs.toy_text.frozen_lake:FrozenLakeEnv',\n",
    "    kwargs={'map_name': '4x4',\n",
    "            'is_slippery': False}) # Note: You have to solve by changing this bool to True.\n",
    "\n",
    "\n",
    "\"\"\" \n",
    "    env: gym.core.Environment\n",
    "    Environment to play on.\n",
    "\n",
    "    env.P: dictionary\n",
    "    It is from gym.core.Environment\n",
    "    P[state][action] is tuples with (probability, nextstate, reward, terminal)\n",
    "    \n",
    "    env.nS: int\n",
    "    count of states \n",
    "    \n",
    "    env.nA: int\n",
    "    count of actions available\n",
    "    \n",
    "    action_space: discrete\n",
    "                LEFT = 0\n",
    "                DOWN = 1\n",
    "                RIGHT = 2\n",
    "                UP = 3\n",
    "    ENVIRONMENT: \n",
    "                \"SFFF\",\n",
    "                \"FHFH\",\n",
    "                \"FFFH\",\n",
    "                \"HFFG\"\n",
    "\"\"\"\n",
    "\n",
    "def print_policy(policy, action_names):\n",
    "    \"\"\" \n",
    "    Print the policy in human-readable format.\n",
    "    \"\"\"\n",
    "    str_policy = policy.astype('str')\n",
    "    for action_num, action_name in action_names.items():\n",
    "        np.place(str_policy, policy == action_num, action_name)\n",
    "\n",
    "    print(str_policy[0:4])\n",
    "    print(str_policy[4:8])\n",
    "    print(str_policy[8:12])\n",
    "    print(str_policy[12:16])\n",
    "    \n",
    "    return str_policy\n",
    "\n",
    "action_names = {LEFT: 'LEFT', RIGHT: 'RIGHT', DOWN: 'DOWN', UP: 'UP'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('D4x4-FrozenLake-v0')\n",
    "grid = 4\n",
    "gamma = 0.9 # Change this to play with it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "########################################################################\n",
    "######################### Policy iteration #############################\n",
    "########################################################################\n",
    "\n",
    "\n",
    "def policy_evaluation(env, gamma, policy, value_func_old, max_iterations=int(1e3), tol=1e-3):\n",
    "    \"\"\"\n",
    "        Evaluate the value of a policy.\n",
    "        See section 4.1 of Reinforcement Learning: An Introduction Second edition by Sutton and Barto\n",
    "    \"\"\"\n",
    "    \n",
    "    value_func_new = np.zeros(env.nS)\n",
    "    value_func_collect = np.zeros((env.nS,max_iterations))\n",
    "\n",
    "    for iteration in range(max_iterations):\n",
    "        \n",
    "        delta=0\n",
    "        \n",
    "        for s in range(env.nS):\n",
    "            value_func_temp = 0\n",
    "            a = policy[s]\n",
    "            \n",
    "            # Using env.P[s][a] implement the V(s) updation given in the pseudocode of Section 4.1\n",
    "            # Note: Since this is a deterministic policy, the summation over\n",
    "            # action space is not required.\n",
    "            \n",
    "            # Enter your code here\n",
    "            \n",
    "            ## Hint: Start by using a \"for\" loop. Access the next states via env.P[s][a]\n",
    "            \n",
    "            # Few code compatibility instructions:\n",
    "            ### Updated V(s) must be in the variable \"value_func_temp\"\n",
    "            ### The older value of V(s) is accessed by the variable \"value_func_old\"\n",
    "            \n",
    "            diff=abs(value_func_old[s]-value_func_temp)\n",
    "            delta=max(delta,diff)\n",
    "            \n",
    "            value_func_new[s]=value_func_temp\n",
    "        \n",
    "        # Stopping criteria: STOP when the sup norm of (V_k-V_{k-1}) is less than some tolerance level(it's 1e-3 here).\n",
    "        if delta<=tol: break\n",
    "        \n",
    "        value_func_old = value_func_new\n",
    "        value_func_collect[:,iteration] = value_func_old\n",
    "\n",
    "    return delta, value_func_new, iteration\n",
    "\n",
    "\n",
    "\n",
    "def policy_improvement(env, gamma, value_func, policy):\n",
    "    \"\"\"\n",
    "      Given a policy and value function, improve the policy.\n",
    "      Returns true if policy is unchanged. Also returns the new policy.\n",
    "      See section 4.2 of Reinforcement Learning: An Introduction Second edition by Sutton and Barto\n",
    "    \"\"\"\n",
    "    value_func_new = np.zeros(env.nS)\n",
    "    policy_stable=True\n",
    "    for s in range(env.nS):\n",
    "        old_action=policy[s]\n",
    "        max_value_func=-1\n",
    "        max_action=-1\n",
    "        for a in range(env.nA):\n",
    "            value_func_temp=0\n",
    "            \n",
    "            # Copy and paste the V(s) updation which you've implemented in the \n",
    "            # policy_evaluation function here (unchanged).\n",
    "            \n",
    "            \n",
    "            if value_func_temp>max_value_func:\n",
    "                max_value_func=value_func_temp\n",
    "                max_action=a\n",
    "        if max_action!=old_action: policy_stable=False\n",
    "        policy[s]=max_action\n",
    "        value_func_new[s]=max_value_func\n",
    "    return policy_stable, policy, value_func_new\n",
    "\n",
    "\n",
    "def policy_iteration(env, gamma, max_iterations=int(1e3), tol=1e-3):\n",
    "    \"\"\"\n",
    "       Runs policy iteration.\n",
    "       Returns optimal policy, value function, number of policy\n",
    "       improvement iterations, and number of value iterations.\n",
    "       See section 4.3 of Reinforcement Learning: An Introduction Second edition by Sutton and Barto\n",
    "    \"\"\"\n",
    "    policy = np.random.randint(4, size=env.nS)\n",
    "    value_func_old = np.random.rand(env.nS)\n",
    "    value_func = np.zeros(env.nS)\n",
    "    value_func_collect = np.zeros((env.nS,max_iterations))\n",
    "    delta_collect = np.zeros(max_iterations)\n",
    "    policy_stable=False\n",
    "    iters=0\n",
    "    eval_iters=0\n",
    "    while not policy_stable:\n",
    "        delta,value_func,iter=policy_evaluation(env,gamma,policy,value_func_old)\n",
    "        delta_collect[iters] = delta\n",
    "        value_func_collect[:,iters] = value_func\n",
    "        eval_iters+=iter\n",
    "        policy_stable,policy,value_func_old=policy_improvement(env,gamma,value_func,policy)\n",
    "        iters+=1\n",
    "\n",
    "    temp=np.repeat(value_func_collect[:,iters-1], iters, axis=0) \n",
    "    temp=temp.reshape(16,iters)\n",
    "    temp1 = abs(value_func_collect[:,0:iters]-temp)\n",
    "    distance_from_converged_value_func = temp1.max(axis=0)\n",
    "\n",
    "    fig = plt.figure()\n",
    "    fig.suptitle(\"Policy iteration convergence plots\", fontsize=16)\n",
    "    ax = plt.subplot(2,1,1)\n",
    "    ax.plot(range(1,iters+1),delta_collect[0:iters])\n",
    "    ax.set_title('||V_PIk-V_PI{k-1}|| plot')\n",
    "    ax.set_yscale('log')\n",
    "\n",
    "\n",
    "    ax = plt.subplot(2,1,2)\n",
    "    ax.plot(range(1,iters+1),distance_from_converged_value_func)\n",
    "    ax.set_title('||V_PIk-V^*|| plot')\n",
    "    ax.set_yscale('log')\n",
    "\n",
    "    return policy, value_func, iters, eval_iters\n",
    "\n",
    "\n",
    "print(\"Doing Policy Iteration\")\n",
    "start_time=time.time()\n",
    "policy, value_func, policy_iters, val_iters= policy_iteration(env,gamma)\n",
    "print(\"Total time taken: \"+str((time.time()-start_time)))\n",
    "print(\"Total Policy Improvement Steps: \"+str(policy_iters))\n",
    "print(\"Total Policy Evaluation Steps: \"+str(val_iters))\n",
    "print(\"Policy:\")\n",
    "policy_str=print_policy(policy,action_names)\n",
    "\n",
    "f, ax = plt.subplots(figsize=(11, 9))\n",
    "cmap = seaborn.diverging_palette(220, 10, as_cmap=True)\n",
    "reshaped=np.reshape(value_func,(grid,grid))\n",
    "seaborn.heatmap(reshaped, cmap=cmap, vmax=1.1,\n",
    "            square=True, xticklabels=grid+1, yticklabels=grid+1,\n",
    "            linewidths=.5, cbar_kws={\"shrink\": .5}, ax=ax, annot=True, fmt=\"f\")\n",
    "counter=0\n",
    "for j in range(0, 4):\n",
    "    for i in range(0, 4):\n",
    "        if policy_str[counter]==\"DOWN\":\n",
    "            plt.text(i+0.5, j+0.7, u'\\u2193', fontsize=12)\n",
    "        elif policy_str[counter]==\"UP\":\n",
    "            plt.text(i+0.5, j+0.7, u'\\u2191', fontsize=12)\n",
    "        elif policy_str[counter]==\"LEFT\":\n",
    "            plt.text(i+0.5, j+0.7, u'\\u2190', fontsize=12)\n",
    "        else:\n",
    "            plt.text(i+0.5, j+0.7, u'\\u2192', fontsize=12)\n",
    "        counter=counter+1\n",
    "        \n",
    "plt.title('Heatmap of policy iteration with value function values and directions')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(5 points) Q1. What do you infer from the convergence plots?\n",
    "\n",
    "(5 points) Q2. What do the numbers in the heatmap mean? (Hint: Check the code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "########################################################################\n",
    "#################### Final policy animation ############################\n",
    "########################################################################\n",
    "\n",
    "flag=input(\"\\nEnter 'Y' if you want to see the final animation of the policy achieved. Else enter something else.\\n\")\n",
    "if flag==\"Y\" or flag==\"y\": print(\"Final Policy Animation\")\n",
    "def run_policy(env,gamma,policy):\n",
    "    initial_state = env.reset()\n",
    "    env.render()\n",
    "    current_state = initial_state\n",
    "    while True:\n",
    "        nextstate, reward, done, debug_info = env.step(policy[current_state])\n",
    "        env.render()\n",
    "\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "        current_state=nextstate\n",
    "        time.sleep(1)\n",
    "\n",
    "if flag==\"Y\" or flag==\"y\": run_policy(env,gamma,policy)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}