01_MDPs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Markov Decision Process (MDP)\n",
    "\n",
    "MDPs provide the mathematical framework to define the idealized Reinforcement Learning problem:\n",
    "\n",
    "<div>\n",
    "  <img src=\"assets/01_MDP_agent_env_diagram.png\">\n",
    "  <br>\n",
    "  <small>Drawing from Sutton and Barto, Reinforcement Learning: An Introduction, 2020</small>\n",
    "</div>\n",
    "\n",
    "**Reinforcement learning** algorithms enable the agent to learn to solve a control task (formulated\n",
    "as an MDP) by trial and error, without any supervision, just interacting with the environment, and\n",
    "receving feedback in the form of positive / negative rewards.\n",
    "\n",
    "## MDP Formal Definition\n",
    "\n",
    "* Set of states $S$\n",
    "* Set of actions $A$\n",
    "* Transition function $P(s'|s, a)$\n",
    "* Reward function $R(s, a, s')$\n",
    "* Initial state $s_0$\n",
    "* Discount factor $\\gamma$\n",
    "* Horizon $H$ (how long we are going to be acting)\n",
    "\n",
    "MDPs are **memoryless** (_Markov property_): we only need the current state to decide what action to\n",
    "take next, and not the history of all the states and actions it took before."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hands-on Example: Grid World\n",
    "\n",
    "<div>\n",
    "  <img src=\"assets/01_MDP_gridworld.png\">\n",
    "  <br>\n",
    "  <small>Picture from the U.C. Berkeley <a href=\"https://inst.eecs.berkeley.edu/~cs188/fa23/\">\n",
    "  CS 188</a> materials.</small>\n",
    "</div>\n",
    "\n",
    "### Let's create a basic _grid_ representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import StrEnum\n",
    "\n",
    "class Cell(StrEnum):\n",
    "    \"\"\"The type of cells in the Grid World.\"\"\"\n",
    "    START  = 'S'\n",
    "    TARGET = 'T'\n",
    "    EMPTY  = 'E'\n",
    "    WALL   = 'W'\n",
    "    BOMB   = 'B'\n",
    "\n",
    "class Grid:\n",
    "    \"\"\"The definition of the grid in the Grid World.\"\"\"\n",
    "\n",
    "    def __init__(self, spec: list[str]):\n",
    "        \"\"\"Creates the grid.\n",
    "\n",
    "        It accepts a list of strings identifying the cell types, for example:\n",
    "        [\n",
    "            'EET',\n",
    "            'EEE',\n",
    "            'SEE'\n",
    "        ]\n",
    "        The start cell is at position (x=0, y=0), the target cell at position (x=2, y=2).\n",
    "        \"\"\"\n",
    "        # TODO: Set the height and the width of the grid. Feel free to assume a well-formed input\n",
    "        #       (i.e., the grid is rectangular and all strings have the same length).\n",
    "        self.height: int = None\n",
    "        self.width:  int = None\n",
    "        # TODO: Store a list of lists of `Cell`s, parsed from the input `spec`.\n",
    "        self.cells:  list[list[Cell]] = None\n",
    "\n",
    "    def __str__(self) -> str:\n",
    "        \"\"\"Returns the string representation of the grid.\n",
    "\n",
    "        For the example above in the __init__ method, the string representation looks like:\n",
    "        E E T\n",
    "        E E E\n",
    "        S E E\n",
    "        \"\"\"\n",
    "        # TODO: Return the string representation.\n",
    "        return '<to implement>'\n",
    "\n",
    "    def __getitem__(self, key) -> Cell:\n",
    "        \"\"\"Returns the cell at position defined by key=(x, y).\"\"\"\n",
    "        # TODO: Return the Cell at position (x, y). Make sure the coordinates match the convention\n",
    "        #       described in the __init__ method!!!\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the implementation!\n",
    "WORLD_GRID = Grid([\n",
    "    'EEET',\n",
    "    'EWEB',\n",
    "    'SEEE',\n",
    "])\n",
    "\n",
    "assert WORLD_GRID[3, 1] == Cell.BOMB\n",
    "assert WORLD_GRID[0, 0] == Cell.START\n",
    "assert WORLD_GRID.height == 3\n",
    "assert WORLD_GRID.width == 4\n",
    "\n",
    "assert str(WORLD_GRID) == \"\"\"\n",
    "E E E T\n",
    "E W E B\n",
    "S E E E\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Policy\n",
    "\n",
    "What does it mean to solve an MDP? It means finding a _policy_ $\\pi$ that maximizes the rewards\n",
    "collected by the agent. $\\pi$ is a function that returns the action the agent should take in the\n",
    "next step given the current state.\n",
    "\n",
    "There are two types of policies:\n",
    "\n",
    " * **Deterministic:** $\\pi(s) = a$\n",
    " \n",
    "   Always returns the same action for a given state.\n",
    " * **Stochastic:** $\\pi(s) = a \\sim P[A|S = s]$\n",
    "\n",
    "   Always samples that action from the same _probability distribution_ for a given state.\n",
    "\n",
    "As an example, let's see how good a random policy is in our grid world!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For reproducibility, all the examples have fixed random seed.\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "def init_random(random_seed=10):\n",
    "    \"\"\"Initializes the random generators used in the code to a predetermined random seed.\"\"\"\n",
    "    random.seed(random_seed)\n",
    "    np.random.seed(random_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import some predefined constants and utility functions for visualization!\n",
    "from util.gridworld import RANDOM_POLICY, GRID_WORLD_MDP, run_simulation\n",
    "\n",
    "init_random()\n",
    "run_simulation(GRID_WORLD_MDP, RANDOM_POLICY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's now create the grid world MDP!\n",
    "\n",
    "Keep in mind that the MDP is defined by: set of states, set of actions, transition and reward\n",
    "functions, and the initial state.\n",
    "\n",
    "Note that the environment is probabilistic: the agent might choose an action, but it will be\n",
    "successful only with a given probability. Put it in another way, there is a non-zero probability to\n",
    "land in a state we did not intended - see the TODOs!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from dataclasses import dataclass\n",
    "from enum import Enum\n",
    "from numpy.testing import assert_almost_equal\n",
    "\n",
    "class Action(Enum):\n",
    "    \"\"\"Enum representing the available actions in the Grid World.\"\"\"\n",
    "    UP    = 1\n",
    "    RIGHT = 2\n",
    "    DOWN  = 3\n",
    "    LEFT  = 4\n",
    "\n",
    "@dataclass(frozen=True)\n",
    "class State:\n",
    "    \"\"\"Class representing a state in the Grid World.\"\"\"\n",
    "    x: int = 0\n",
    "    y: int = 0\n",
    "\n",
    "    def pos(self) -> tuple[int]:\n",
    "        return (self.x, self.y)\n",
    "\n",
    "class GridMDP:\n",
    "    \"\"\"Class modeling the Grid World Markov Decision Process (MDP).\n",
    "\n",
    "    It defines the transition and reward functions, as well as states and gamma factor.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, grid: Grid, start = State(), gamma = 1.0):\n",
    "        \"\"\"Initializes the GridMDP.\"\"\"\n",
    "        # TODO: Store the grid, start state, and gamma factor.\n",
    "        self.grid = None\n",
    "        self.start = None\n",
    "        self.gamma = None\n",
    "        # TODO: Enumerate all available actions.\n",
    "        self.all_actions = None\n",
    "        # TODO: Enumerate all possible states. We could filter out the non-reachable states here,\n",
    "        #       but let's keep things simple and enumerate all of them in the grid.\n",
    "        self.all_states = None\n",
    "\n",
    "    def is_terminal(self, state: State) -> bool:\n",
    "        \"\"\"Determines if the `state` is a terminal state.\"\"\"\n",
    "        # TODO: Return whether the state is a terminal state. Hint: TARGET and BOMB are the only\n",
    "        #       terminal states in the Grid World.\n",
    "        return None\n",
    "\n",
    "    def is_reachable(self, state: State) -> bool:\n",
    "        \"\"\"Determines if the `state` is reachable.\"\"\"\n",
    "        # TODO: Return whether the state is reachable. Hint: a state is reachable if it is inside\n",
    "        #       the grid, and it is not a WALL.\n",
    "        return None\n",
    "\n",
    "    def reward(self, state: State, action: Action, next_state: State) -> float:\n",
    "        \"\"\"The reward function of the Grid World MDP.\"\"\"\n",
    "        # TODO: Compute and return the reward for the (state, action, next_state) tuple. We'll only\n",
    "        #       use next_state. If the next_state is TARGET, return 1.0 * gamma. If is is a BOMB,\n",
    "        #       return -1.0 * gamma. Else return zero. \n",
    "        return None\n",
    "\n",
    "    def transition(self, state: State, action: Action, noise = 0.0) -> dict[State, float]:\n",
    "        \"\"\"The transition function of the Grid World MDP.\"\"\"\n",
    "        if not self.is_reachable(state) or self.is_terminal(state):\n",
    "            return {}\n",
    "\n",
    "        def landing(candidate: State) -> State:\n",
    "            \"\"\"Returns the state that we will end up, when attempting to reach another state.\n",
    "\n",
    "            Basically, if `candidate` is reachable it returns `candidate`, otherwise we stay put.\n",
    "            \"\"\"\n",
    "            return candidate if self.is_reachable(candidate) else state\n",
    "\n",
    "        # TODO: Compute the action expected \"landings\" for all the actions.\n",
    "        action_landings = {}\n",
    "        # TODO: Compute the direction the agent might go by mistake (instead of the direction of the\n",
    "        #       chosen action) given the probabilistic Grid World MDP dynamics.\n",
    "        #       Hint: an action can go wrong only in the orthogonal directions (e.g., the agent\n",
    "        #       tries to go left, but it ends up going up or down; it cannot go right if it attemps\n",
    "        #       to go left).\n",
    "        mistaken_directions = None\n",
    "\n",
    "        # TODO: Compute the next state using the action_landings.\n",
    "        next_state = None\n",
    "        # TODO: Compute the list of possible \"mistake\" states given the mistaken_directions.\n",
    "        mistake_states = None\n",
    "\n",
    "        # Compute the transition probabilities. Probabilities are zero in every state, besides the\n",
    "        # next expected state from the chosen action - which has probability (1.0 -noise) - and the\n",
    "        # potentially \"mistake\" states which all have probability (noise / n_possible_mistakes). \n",
    "        probs = defaultdict(lambda: 0.0)\n",
    "        # TODO: Compute the probability of the next expected state.\n",
    "        probs[next_state] = None\n",
    "        for m in mistake_states:\n",
    "            # TODO: Compute the probability of the mistaken states.\n",
    "            probs[m] += None\n",
    "\n",
    "        # A couple of assertions to verify the correctness of the probability computation.\n",
    "        assert sum(probs.values()) <= 1.0\n",
    "        assert_almost_equal(sum(probs.values()), 1.0)\n",
    "\n",
    "        return probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the implementation!\n",
    "from numpy.testing import assert_almost_equal\n",
    "\n",
    "init_random()\n",
    "\n",
    "GRID_MDP = GridMDP(WORLD_GRID, gamma=0.9)\n",
    "\n",
    "# Test all states\n",
    "all_states = GRID_MDP.all_states\n",
    "assert len(all_states) == 4 * 3\n",
    "assert GRID_MDP.all_actions == list(Action)\n",
    "\n",
    "# Test terminal and reachable states\n",
    "for s in all_states:\n",
    "    if s == State(3, 2) or s == State(3, 1):\n",
    "        assert GRID_MDP.is_terminal(s)\n",
    "        assert GRID_MDP.is_reachable(s)\n",
    "    elif s == State(1, 1):\n",
    "        assert not GRID_MDP.is_reachable(s)\n",
    "    else:\n",
    "        assert not GRID_MDP.is_terminal(s)\n",
    "        assert GRID_MDP.is_reachable(s)\n",
    "\n",
    "# Test rewards\n",
    "assert GRID_MDP.reward(State(0, 2), Action.DOWN, State(0, 1)) == 0.0\n",
    "assert GRID_MDP.reward(State(2, 1), Action.RIGHT, State(3, 1)) == -0.9\n",
    "assert GRID_MDP.reward(State(2, 2), Action.RIGHT, State(3, 2)) == 0.9\n",
    "\n",
    "# Test transitions\n",
    "p = GRID_MDP.transition(State(1, 2), Action.RIGHT, noise=0.2)\n",
    "assert_almost_equal(p[State(1, 2)], 0.2)\n",
    "assert_almost_equal(p[State(2, 2)], 0.8)\n",
    "\n",
    "p = GRID_MDP.transition(State(2, 0), Action.UP)\n",
    "assert_almost_equal(p[State(2, 1)], 1.0)\n",
    "assert_almost_equal(p[State(2, 0)], 0.0)\n",
    "assert_almost_equal(p[State(1, 0)], 0.0)\n",
    "assert_almost_equal(p[State(3, 0)], 0.0)\n",
    "assert State(1, 1) not in p\n",
    "\n",
    "p = GRID_MDP.transition(State(1, 0), Action.UP) # Hitting wall, for sure stays.\n",
    "assert_almost_equal(p[State(1, 0)], 1.0)\n",
    "assert_almost_equal(p[State(2, 0)], 0.0)\n",
    "assert_almost_equal(p[State(0, 0)], 0.0)\n",
    "\n",
    "p = GRID_MDP.transition(State(2, 0), Action.UP, noise=0.2)\n",
    "assert_almost_equal(p[State(2, 1)], 0.8)\n",
    "assert_almost_equal(p[State(1, 0)], 0.1)\n",
    "assert_almost_equal(p[State(3, 0)], 0.1)\n",
    "assert_almost_equal(p[State(2, 0)], 0.0)\n",
    "\n",
    "p = GRID_MDP.transition(State(0, 0), Action.DOWN, noise=0.2)\n",
    "assert_almost_equal(p[State(0, 0)], 0.9)\n",
    "assert_almost_equal(p[State(1, 0)], 0.1)\n",
    "assert_almost_equal(p[State(0, 1)], 0.0)\n",
    "\n",
    "p = GRID_MDP.transition(State(0, 0), Action.UP, noise=0.2)\n",
    "assert_almost_equal(p[State(0, 0)], 0.1)\n",
    "assert_almost_equal(p[State(1, 0)], 0.1)\n",
    "assert_almost_equal(p[State(0, 1)], 0.8)\n",
    "\n",
    "p = GRID_MDP.transition(State(2, 1), Action.LEFT, noise=0.2)\n",
    "assert_almost_equal(p[State(2, 1)], 0.8)\n",
    "assert_almost_equal(p[State(2, 2)], 0.1)\n",
    "assert_almost_equal(p[State(2, 0)], 0.1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Goal\n",
    "\n",
    "RL is based on the _reward hypothesis_: all goals can be formulated as the maximization of the agent\n",
    "(discounted) _cumulative_ reward:\n",
    "\n",
    "$$\n",
    "G_t = R_{t+1} + \\gamma R_{t+2} + \\gamma^2 R_{t+3} + ... = \\sum_{k=0}^{H} \\gamma^k R_{t + k + 1}\n",
    "$$\n",
    "\n",
    "Why discounts? To value of short-term rewards vs. future ones (e.g., money / interests).\n",
    "\n",
    "More formally, the agent goal is to find a policy (usually indicated with $\\pi$) that maximizes such\n",
    "expected return:\n",
    "\n",
    "$$\n",
    "Goal: \\max_{\\pi} \\mathbb{E}[ G_t | \\pi ] =\n",
    "      \\max_{\\pi} \\mathbb{E}[ \\sum_{t=0}^{H} \\gamma^k R(S_t, A_t, S_{t+1}) | \\pi ]\n",
    "$$\n",
    "\n",
    "Such policy is called the _optimal_ policy $\\pi^*$. Note: there can be more than one!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exact Solution Methods\n",
    "\n",
    "When the MDP is small enough and we have access to the full dynamics (or transition model) of the\n",
    "environment, we can solve the MDP problem _exactly_ using dynamic programming.\n",
    "\n",
    "### Optimal Value Functions\n",
    "\n",
    "**Intuition:** A _value function_ assigns a numerical value to each state, representing how \"good\"\n",
    "such state is for the agent. For now, let's assume we already know the optimal policy.\n",
    "\n",
    "Two types of (optimal) value functions:\n",
    "\n",
    " * $V^*(s)$ is the best return we can get starting from state $s$ and then acting optimally.\n",
    " * $Q^*(s, a)$ is the best return we can get starting from state $s$ taking action $a$ and then\n",
    "   acting optimally.\n",
    "\n",
    "Generally, it is true that: $V^*(s) = \\max_{a} Q^*(s, a)$\n",
    "\n",
    "### Value Iteration\n",
    "\n",
    "The \"value iteration\" algorithm is based on the following recursive relation (or Bellman equation):\n",
    "\n",
    "$$\n",
    "Q^*(s, a) = \\sum_{s'} P(s'|s, a)[R(s,a,s') + \\gamma V^*(s')]\n",
    "$$\n",
    "\n",
    "Basically, the value of a certain state `s` is given by the sum (for all possible next states) of\n",
    "the probability of transitioning to that next state `s'` multiplied by the reward of that transition\n",
    "_plus_ the discounted value of such next state (hence, the recursive term).\n",
    "\n",
    "For the mathematical derivation, checked the linked materials in the Intro :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "class QTable:\n",
    "    \"\"\"Table storing Q, i.e., the state-action values.\n",
    "\n",
    "    This is not an optimal implementation, but hopefully useful for learning purposes.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, states: list[State], actions: list[Action]):\n",
    "        \"\"\"Initializes the QTable given all states and actions of the MDP.\"\"\"\n",
    "        self.states = states\n",
    "        self.actions = actions\n",
    "        self.nA = len(actions)\n",
    "        # Why a dict of dict instead of a single dict indexed by tuple (State, Action)? Just for\n",
    "        # convenience when looking up all action values for a state.\n",
    "        self.table: dict[State, dict[Action, float]] = \\\n",
    "            { s : { a : 0.0 for a in actions } for s in states }\n",
    "\n",
    "    def __getitem__(self, key: tuple[State, Action]) -> float:\n",
    "        \"\"\"Returns the Q value for a tuple of (state, action).\"\"\"\n",
    "        # TODO: Return the corresponding value in the internal table.\n",
    "        return None\n",
    "\n",
    "    def __setitem__(self, key: tuple[State, Action], value: float):\n",
    "        \"\"\"Returns the Q value for a tuple of (state, action).\"\"\"\n",
    "        # TODO: Set the value in the internal table.\n",
    "        pass\n",
    "\n",
    "    def value(self, state: State) -> float:\n",
    "        \"\"\"Returns the value of a state.\"\"\"\n",
    "        # TODO: Return the state value. Hint: the value of a state is the max Q value across all\n",
    "        #       actions for that state.\n",
    "        return None\n",
    "\n",
    "    def best_action(self, state: State) -> Action:\n",
    "        \"\"\"Returns the best action for a certain state.\"\"\"\n",
    "        best_action = None\n",
    "        best_v = float('-inf')\n",
    "        actions = list(self.table[state].keys())\n",
    "        # TODO: Use random.shuffle to shuffle the actions (in case multiple have the same value).\n",
    "        for a in actions:\n",
    "            # TODO: Search for the best value.\n",
    "            pass\n",
    "        return best_action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test our implementation!\n",
    "init_random()\n",
    "\n",
    "state_0 = State(0, 0)\n",
    "state_1 = State(10, 10)\n",
    "state_2 = State(11, 11)\n",
    "state_3 = State(11, 10)\n",
    "qtable = QTable([state_0, state_1, state_2, state_3], list(Action))\n",
    "\n",
    "# props\n",
    "assert qtable.states == [state_0, state_1, state_2, state_3]\n",
    "assert qtable.actions == list(Action)\n",
    "assert qtable.nA == len(list(Action))\n",
    "\n",
    "# init get\n",
    "assert qtable[(State(10, 10), Action.DOWN)] == 0.0\n",
    "\n",
    "# set\n",
    "qtable[state_1, Action.DOWN] = 0.5\n",
    "qtable[state_1, Action.UP]   = 1.5\n",
    "qtable[state_2, Action.LEFT] = 2.5\n",
    "assert qtable[state_1, Action.DOWN]  == 0.5\n",
    "assert qtable[state_1, Action.UP]    == 1.5\n",
    "assert qtable[state_2, Action.LEFT]  == 2.5\n",
    "assert qtable[state_2, Action.RIGHT] == 0.0\n",
    "\n",
    "# value\n",
    "assert qtable.value(state_1) == 1.5\n",
    "assert qtable.value(state_2) == 2.5\n",
    "assert qtable.value(state_3) == 0.0\n",
    "\n",
    "# best action\n",
    "assert qtable.best_action(state_1) == Action.UP\n",
    "assert qtable.best_action(state_2) == Action.LEFT\n",
    "# let's hope large numbers don't fail us :)\n",
    "many_actions = [qtable.best_action(state_0) for _ in range(5_000)]\n",
    "assert len(dict.fromkeys(many_actions)) == 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def value_iteration(mdp: GridMDP, noise = 0.0, n_iterations = 100) -> QTable:\n",
    "    \"\"\"Runs the value iteration algorithm for the Grid World MDP\"\"\"\n",
    "    # TODO: Initialize the QTable with the states and actions coming from the MDP.\n",
    "    qtable = None\n",
    "\n",
    "    for _ in range(0, n_iterations):\n",
    "        # TODO: Create a new QTable to store the updated values.\n",
    "        new_qtable = None\n",
    "        # TODO: Loop over all states, over all actions, get the transition probabilities, get the\n",
    "        #       reward via the MDP reward function, and update the new_qtable according to the\n",
    "        #       update-rule defined above (i.e., prob * (reward + gamma * state_value))\n",
    "        \n",
    "        # Swap the qtable with the updated new_qtable.\n",
    "        qtable = new_qtable\n",
    "\n",
    "    return qtable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test our implementation!\n",
    "init_random()\n",
    "\n",
    "qtable = value_iteration(GRID_MDP, noise = 0.2)\n",
    "\n",
    "assert f'{qtable.value(State(0, 0)):.2f}' == '0.49'\n",
    "assert f'{qtable.value(State(1, 0)):.2f}' == '0.43'\n",
    "assert f'{qtable.value(State(2, 0)):.2f}' == '0.48'\n",
    "assert f'{qtable.value(State(3, 0)):.2f}' == '0.28'\n",
    "assert f'{qtable.value(State(0, 1)):.2f}' == '0.57'\n",
    "assert f'{qtable.value(State(1, 1)):.2f}' == '0.00'\n",
    "assert f'{qtable.value(State(2, 1)):.2f}' == '0.57'\n",
    "assert f'{qtable.value(State(3, 1)):.2f}' == '0.00'\n",
    "assert f'{qtable.value(State(0, 2)):.2f}' == '0.64'\n",
    "assert f'{qtable.value(State(1, 2)):.2f}' == '0.74'\n",
    "assert f'{qtable.value(State(2, 2)):.2f}' == '0.85'\n",
    "assert f'{qtable.value(State(3, 2)):.2f}' == '0.00'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize our value function!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using provided utilities, we can print the state values of the Grid World!\n",
    "from util.gridworld import plot_grid\n",
    "\n",
    "_, _ = plot_grid(GRID_MDP.grid, qtable)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Notes\n",
    "\n",
    "**Vectorized implementation**: This is clearly not an optimal implementation, but it is good for\n",
    "understanding the algorithm. For an optimized algorithm, we should vectorize states, actions, and\n",
    "use a vectorized implementation. Something along the lines of:\n",
    "\n",
    "```python\n",
    "    new_value = np.max(np.sum(transitions * (rewards + gamma * old_value), axis=2), axis=1)\n",
    "```\n",
    "\n",
    "**Convergence**: Value iteration (as well as policy iteration, see below) are proven to converge.\n",
    "Check out the materials linked in the Intro in case you are interested.\n",
    "\n",
    "**Policy Iteration**: Another exect solution algorithm (one \"policy evaluation\" step to find state\n",
    "values, followed by \"policy improvement\"). It is usually faster (see linked resources)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimal Policy\n",
    "\n",
    "Got the optimal value function, and now? Well, we are in luck: an optimal policy is the one that chooses the action with higher value!\n",
    "\n",
    "$$\n",
    "\\pi^*(s) = \\argmax_a Q^*(s, a)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_policy = lambda s: qtable.best_action(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see how an optimal agent behaves!\n",
    "init_random()\n",
    "run_simulation(GRID_MDP, optimal_policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aren't We Done?! :)\n",
    "\n",
    "Fortunately for our own amusement, we are not! There are two major very practical limitations:\n",
    "\n",
    " 1. We rarely (almost never!) have access to the MDP transition model.\n",
    " 2. The state and action space are generally large and not tractable this way.\n",
    "\n",
    "The first point is what classic reinforcement learning algorithms address: enable the agent to find\n",
    "an optimal policy just from interacting with (i.e., sampling) the environment. The second category\n",
    "is addressed by deep reinforcement learning! Keep going and have fun!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Grid World Environment\n",
    "\n",
    "The _environment_ is what the agent will use in reinforcement learning algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class GridEnv:\n",
    "    \"\"\"A reinforcement learning environment to interact with the Grid World and sample from it.\"\"\"\n",
    "\n",
    "    def __init__(self, mdp: GridMDP):\n",
    "        \"\"\"Initializes the GridEnv.\"\"\"\n",
    "        self.mdp = mdp\n",
    "        self.state = mdp.start\n",
    "        self.terminated = False\n",
    "\n",
    "    def reset(self) -> State:\n",
    "        \"\"\"Resets the state to the start state, not terminated, and returns the new start state.\"\"\"\n",
    "        self.state = self.mdp.start\n",
    "        self.terminated = False\n",
    "        return self.state\n",
    "\n",
    "    def step(self, action: Action) -> tuple[State, float, bool]:\n",
    "        \"\"\"Performs a step in the environment given the selected `action` by the agent.\n",
    "\n",
    "        Returns a tuple of (next_state, reward, is_done), with the next state after the action has\n",
    "        been taked, the associated reward, and whether the episode terminated.\n",
    "        \"\"\"\n",
    "        if self.terminated:\n",
    "            raise Exception('Environment episode completed, please call reset.')\n",
    "\n",
    "        # TODO: Compute the state probabilities for the action, sample the probabilities to select\n",
    "        #       the next state, calculate the reward, determine if the next_state is terminal, and\n",
    "        #       return the tuple (next_state, reward, done).\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your implementation!\n",
    "init_random()\n",
    "\n",
    "GRID_ENV = GridEnv(GridMDP(WORLD_GRID))\n",
    "\n",
    "assert not GRID_ENV.terminated\n",
    "\n",
    "assert GRID_ENV.step(Action.UP) == (State(0, 1), 0.0, False)\n",
    "assert GRID_ENV.state == State(0, 1)\n",
    "\n",
    "assert GRID_ENV.step(Action.UP) == (State(0, 2), 0.0, False)\n",
    "assert GRID_ENV.state == State(0, 2)\n",
    "\n",
    "GRID_ENV.terminated = True\n",
    "assert GRID_ENV.reset() == State()\n",
    "assert GRID_ENV.state == State()\n",
    "assert GRID_ENV.terminated == False\n",
    "\n",
    "GRID_ENV.step(Action.UP)\n",
    "GRID_ENV.step(Action.UP)\n",
    "GRID_ENV.step(Action.RIGHT)\n",
    "\n",
    "assert GRID_ENV.step(Action.RIGHT) == (State(2, 2), 0.0, False)\n",
    "assert GRID_ENV.step(Action.RIGHT) == (State(3, 2), 1.0, True)\n",
    "assert GRID_ENV.terminated\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Appendix\n",
    "\n",
    "### Noise and Discount Tradeoff\n",
    "\n",
    "Consider the cliff-world grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from util.gridworld import Grid\n",
    "\n",
    "# N has -10 return, G has +10 return.\n",
    "CLIFF_WORLD_GRID = Grid([\n",
    "    'EEEEE',\n",
    "    'EWEEE',\n",
    "    'EWTWG',\n",
    "    'SEEEE',\n",
    "    'NNNNN',\n",
    "])\n",
    "\n",
    "plot_grid(CLIFF_WORLD_GRID, agent_pos=(0, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try to solve the environment with different values of $\\gamma$ (e.g., `0.1` and `0.99`) and _noise_\n",
    "(e.g., `0` and `0.5`). When does the agent prefers the shortest path? What about the closer exit?\n",
    "When does it risk the cliff?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extensions\n",
    "\n",
    "There are extensions to MDP such as POMDP (Partially Observable MDP) and HMM (Hidden Markov Models)\n",
    "that are interesting areas to explore and learn about as well!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "drlzh",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}