ROS_Core/src/Labs/Lab3/pomdp.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "816729b5-ad6e-4131-934e-7d84d28a2f13",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Lab 3: MDP and QMDP\n",
    "This lab aims to show you how to formulate a problem as an MDP (POMDP) and solve it using value iteration, policy iteration, and approximated solution for POMDP such as QMDP.\n",
    "\n",
    "## 1. MDP\n",
    "### Task 1.1: MDP formulation\n",
    "We provide you with a class `MDP()` to help you easily formulate your MDP problem. Refer to your lab handout for instructions on how to use the class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27599c2e-1945-4e5f-b275-716b66568bfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import os\n",
    "from mdp import MDP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a781589c-8a5a-4e9b-bdd5-26870272f629",
   "metadata": {},
   "source": [
    "Test out the class with the simple two-state problem in the handout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f364f75c-101a-4bde-b63b-fce52e49c25f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoStateMDP(MDP):\n",
    "    def __init__(self):\n",
    "        self.states = [[\"s1\", \"s2\"]]\n",
    "        self.actions = [\"a0\", \"a1\"]\n",
    "        self.gam = 0.9\n",
    "    \n",
    "        # call the parent class\n",
    "        # notice that the state is a list of state variables\n",
    "        super().__init__(\n",
    "            states=self.states, actions=self.actions)\n",
    "        self.populate_data()\n",
    "    \n",
    "    def populate_data(self):\n",
    "        # add all routes from s1\n",
    "        self.add_route([\"s1\"],\"a0\",[\"s1\"])\n",
    "        self.add_route([\"s1\"],\"a1\",[\"s2\"])\n",
    "        # add all routes from s2\n",
    "        self.add_route([\"s2\"],\"a0\",[\"s2\"])\n",
    "        self.add_route([\"s2\"],\"a1\",[\"s2\"])\n",
    "        \n",
    "        # let's populate the reward, assuming r>0 is 0.5\n",
    "        for a in self.a:\n",
    "            self.add_reward([\"s1\"], a, 0.5)\n",
    "            self.add_reward([\"s2\"], a, 1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1887459-28a7-49c9-894a-6cfe600b4540",
   "metadata": {},
   "outputs": [],
   "source": [
    "twoStateMDP = TwoStateMDP()\n",
    "assert twoStateMDP.get_index([\"s1\"]) == 0, \"Something is wrong\"\n",
    "assert twoStateMDP.get_state(0) == [0], \"Something is wrong\"\n",
    "assert twoStateMDP.get_real_state_value(0) == ['s1'], \"Something is wrong\"\n",
    "print(\"Everything is correct!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "703348b5-1d8b-4dd4-b025-8eb8cf89472d",
   "metadata": {},
   "source": [
    "Now let's use the `MDP()` class to formulate our T-intersection problem.\n",
    "\n",
    "**Task 1.1**: In `populate_data()`, all the probability value for each `self.add_route()` command is missing (denotes `MISSING_VALUE`). Compute the state transition matrix of the T-intersection MDP and fill in the missing `p` values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4c45cf1-92c5-4b0f-8adc-19787e2ebf77",
   "metadata": {},
   "outputs": [],
   "source": [
    "### TASK 1.1: Fill in MISSING_VALUE with correct p values\n",
    "class TIntersection(MDP):\n",
    "    def __init__(self, reward={\"forward\": -1.0, \"stop\": -5, \"collision\": -10, \"goal\": 5.0}):\n",
    "        self.states = [\n",
    "            [\"ego_{}\".format(x) for x in list(range(1, 6))],\n",
    "            [\"car_{}\".format(x) for x in list(range(1, 6))]\n",
    "        ]\n",
    "        self.actions = [\"forward\", \"stop\"]\n",
    "        self.gam = 0.9\n",
    "        super().__init__(\n",
    "            states=self.states, actions=self.actions, method=\"add\")\n",
    "        self.reward = reward\n",
    "        self.populate_data()\n",
    "    \n",
    "    def populate_data(self):\n",
    "        for i in range(1, 6):\n",
    "            for j in range(1, 6):\n",
    "                # add route for forward action\n",
    "                # (i, j) --> (i+1, j)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+1, 5)), \"car_{}\".format(j)], MISSING_VALUE)\n",
    "                # (i, j) --> (i+1, j+1)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+1, 5)), \"car_{}\".format(min(j+1, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i+1, j+2)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+1, 5)), \"car_{}\".format(min(j+2, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i+2, j)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+2, 5)), \"car_{}\".format(j)], MISSING_VALUE)\n",
    "                # (i, j) --> (i+2, j+1)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+2, 5)), \"car_{}\".format(min(j+1, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i+2, j+2)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+2, 5)), \"car_{}\".format(min(j+2, 5))], MISSING_VALUE)\n",
    "\n",
    "                # add route for stop action\n",
    "                # (i, j) --> (i, j+2)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", [\"ego_{}\".format(i), \"car_{}\".format(min(j+2, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i, j+1)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", [\"ego_{}\".format(i), \"car_{}\".format(min(j+1, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i, j)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", [\"ego_{}\".format(i), \"car_{}\".format(j)], MISSING_VALUE)\n",
    "\n",
    "                self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", self.reward[\"forward\"])\n",
    "                self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", self.reward[\"stop\"])\n",
    "                        \n",
    "                # check for collision\n",
    "                if (i in [4, 5]) and (j in [4, 5]):\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", self.reward[\"collision\"])\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", self.reward[\"collision\"])\n",
    "                elif (i == 5) or (j == 5):\n",
    "                    # reaching goal\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", self.reward[\"goal\"])\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", self.reward[\"goal\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea65f5e7-9958-4399-9656-0e9d9bce091a",
   "metadata": {},
   "outputs": [],
   "source": [
    "### TEST YOUR CODE\n",
    "tinter = TIntersection()\n",
    "init_state = [\"ego_1\", \"car_1\"]\n",
    "init_state_index = tinter.get_index(init_state)\n",
    "state_transition_matrix = tinter.P[:, 0, init_state_index]\n",
    "next_state_index = np.where(state_transition_matrix > 0.0)\n",
    "next_state_p = state_transition_matrix[next_state_index]\n",
    "assert np.array_equal(next_state_index[0], np.array([1, 2, 6, 7, 11, 12]))\n",
    "assert np.array_equal(next_state_p, np.array([0.64, 0.16, 0.128, 0.032, 0.032, 0.008]))\n",
    "\n",
    "print(\"Correct!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e45a7cbc-3334-4b9b-bd50-6c561972703f",
   "metadata": {},
   "source": [
    "### Task 1.2: MDP value iteration and policy iteration\n",
    "Now we will write the value iteration function and policy iteration function or an arbitrary MDP that inherits the structure of our MDP class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2220b35-921f-4646-a773-fec781aedf9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "### TASK 1.2: Write value iteration and policy iteration\n",
    "def value_iteration(threshold = .001, mdp:MDP=None):\n",
    "    if mdp is None:\n",
    "        raise ValueError(\"MDP cannot be None\")\n",
    "    numa, nums, R, P = mdp.get_mdp()\n",
    "    V_star = np.zeros(nums)\n",
    "    pi_star = np.zeros(nums)\n",
    "    \n",
    "    ####\n",
    "    ## YOUR CODE HERE\n",
    "    raise NotImplementedError(\"You have not written Value Iteration\")\n",
    "    ####\n",
    "\n",
    "    return V_star, pi_star\n",
    "\n",
    "def policy_eval(policy, threshold = .001, mdp:MDP=None):\n",
    "    if mdp is None:\n",
    "        raise ValueError(\"MDP cannot be None\")\n",
    "    numa, nums, R, P = mdp.get_mdp()\n",
    "    V = np.zeros(nums)\n",
    "    \n",
    "    ####\n",
    "    ## YOUR CODE HERE\n",
    "    raise NotImplementedError(\"You have not written Policy Evaluation\")\n",
    "    ####\n",
    "    \n",
    "    return V\n",
    "\n",
    "def policy_iteration(threshold = .001, mdp:MDP=None):\n",
    "    if mdp is None:\n",
    "        raise ValueError(\"MDP cannot be None\")\n",
    "    numa, nums, R, P = mdp.get_mdp()\n",
    "    # initialize a random policy with length nums and action randomly assigned from numa\n",
    "    pi_star = np.random.randint(0, numa, nums)\n",
    "    V_star = policy_eval(pi_star, mdp=mdp)\n",
    "\n",
    "    ####\n",
    "    ## YOUR CODE HERE\n",
    "    raise NotImplementedError(\"You have not written Policy Iteration\")\n",
    "    ####\n",
    "\n",
    "    return V_star, pi_star"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "695ec1a3-ea3d-4348-ab62-909ba1fb3997",
   "metadata": {},
   "source": [
    "Let's test our value iteration and policy iteration on the `TwoStateMDP`. The following is the close-form value function for this simple MDP:\n",
    "\n",
    "$$\n",
    "V(s_1) = \\frac{r + \\gamma}{1 - \\gamma} \\;\\;,\\;\\; V(s_2) = \\frac{1 + r}{1 - \\gamma}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72248dd1-1c78-470c-81f5-4e1e244bf51e",
   "metadata": {},
   "outputs": [],
   "source": [
    "### TEST YOUR CODE\n",
    "V_star, pi_star_value = value_iteration(mdp=twoStateMDP, threshold=1e-10)\n",
    "r = 0.5\n",
    "V_calc = [(r + twoStateMDP.gam)/(1-twoStateMDP.gam), (1 + r)/(1-twoStateMDP.gam)]\n",
    "assert np.max(abs(V_star - V_calc)) < 1e-2, \"Value iteration is incorrect\"\n",
    "print(\"Value iteration is correct!\")\n",
    "\n",
    "V_star, pi_star_policy = policy_iteration(mdp=twoStateMDP, threshold=1e-10)\n",
    "assert np.max(abs(V_star - V_calc)) < 1e-2, \"Policy iteration is incorrect\"\n",
    "print(\"Policy iteration is correct!\")\n",
    "\n",
    "assert np.array_equal(pi_star_policy, pi_star_value), \"Policy and value iteration give different pi star\"\n",
    "print(\"Policy learned by value iteration: {}\".format(pi_star_value))\n",
    "print(\"Policy learned by policy iteration: {}\".format(pi_star_policy))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98ae34cf-f9d2-43b6-b727-6ca0829b7d04",
   "metadata": {},
   "source": [
    "Let's now test the value iteration and policy iteration on our T-intersection MDP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8c3a903-8c39-4e98-9086-655d3c666b28",
   "metadata": {},
   "outputs": [],
   "source": [
    "V_star, pi_star = value_iteration(mdp=tinter, threshold=1e-5)\n",
    "print(\"V*: {}\".format(V_star))\n",
    "print(\"pi*: {}\".format(pi_star))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7df012f7",
   "metadata": {},
   "source": [
    "### Task 1.3: Simulate your computed $\\pi^*$\n",
    "\n",
    "We provide you with a class `TintersectionVisualizer()` to visualize your MDP. Simply call the following function to plot the current state:\n",
    "```python\n",
    "# initialize the visualizer\n",
    "vis = TIntersectionVisualizer()\n",
    "# define the state\n",
    "state = [\"ego_1\", \"car_1\"]\n",
    "# visualize the state\n",
    "vis.plot(state)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "805792f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from visualizer import TIntersectionVisualizer\n",
    "vis = TIntersectionVisualizer()\n",
    "state = [\"ego_1\", \"car_1\"]\n",
    "vis.plot(state)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "98e1ead6",
   "metadata": {},
   "source": [
    "Do the following task:\n",
    "1. Choose an initial state\n",
    "2. From the computed $V^*(x), \\pi^*(x)$, iterate from the chosen initial state gradually until you reach the terminal state, i.e. `[\"ego_5\", \"car_5\"]`.\n",
    "3. While iterating, keep track of your state-action pair $(s, a)$. Print out all the state-action pairs that your computed $V^*(x), \\pi^*(x)$ navigate you to.\n",
    "\n",
    "*Hint 1*: For each state and action pair $(s, a)$, the matrix P will show you the probability of $s'$ by calling `P[:, state_index, action_index]`\n",
    "\n",
    "*Hint 2*: Assuming that you have a list of states `S` and the probability distribution `p` of the states in `S`, you can use `numpy.random.choice(S, p=p)` to sample the next state, given `S` and `p`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82336901",
   "metadata": {},
   "outputs": [],
   "source": [
    "## YOUR CODE HERE\n",
    "def get_solution(initial_state):\n",
    "    done = False\n",
    "    state = initial_state \n",
    "    state_list = [state]\n",
    "    print(\"State: {}\\tAction: {}\".format(state, None))\n",
    "\n",
    "    while not done:\n",
    "        ## YOUR CODE HERE\n",
    "        # 1. Get index of current state\n",
    "        # 2. Get next action based on pi_star and state index\n",
    "        # 3. Get indices of all the possible next states and their transition probabilities\n",
    "        # 4. Sample next state index, given the transition probabilities\n",
    "        # 5. Append new state (readable form, e.g. [\"ego_1\", \"car_2\"]) to state_list\n",
    "\n",
    "        if state[0] == \"ego_5\" or state[1] == \"car_5\":\n",
    "            done = True\n",
    "\n",
    "    return state_list\n",
    "\n",
    "initial_state = ... # choose an initial condition\n",
    "state_list = get_solution(initial_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fce6c6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import imageio\n",
    "from IPython.display import Image\n",
    "\n",
    "folder = \"figures\"\n",
    "sub_folder = \"mdp\"\n",
    "\n",
    "fig_prog_folder = os.path.join(folder, sub_folder)\n",
    "if os.path.exists(fig_prog_folder):\n",
    "    print(\"WARNING: Path {} exists, GIF result might be affected with old data\".format(fig_prog_folder))\n",
    "os.makedirs(fig_prog_folder, exist_ok=True)\n",
    "\n",
    "for i, state in enumerate(state_list):\n",
    "    vis.plot(state)\n",
    "    plt.savefig(os.path.join(fig_prog_folder, \"{}.png\".format(i)), dpi=200)\n",
    "    plt.clf()\n",
    "\n",
    "gif_path = os.path.join(fig_prog_folder, 'result.gif')\n",
    "length = len([i for i in os.listdir(os.path.join(fig_prog_folder)) if \".png\" in i])\n",
    "\n",
    "with imageio.get_writer(gif_path, mode='I') as writer:\n",
    "    for i in range(length):\n",
    "        print(i, end='\\r')\n",
    "        filename = os.path.join(fig_prog_folder, str(i)+\".png\")\n",
    "        image = imageio.imread(filename)\n",
    "        writer.append_data(image)\n",
    "Image(open(gif_path,'rb').read(), width=400)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c6c89ab7",
   "metadata": {},
   "source": [
    "Let's play around with this a bit. Let's change the default reward function of our MDP.\n",
    "\n",
    "The class `TIntersection()` takes in a dictionary with reward information of the following structure:\n",
    "\n",
    "```python\n",
    "TIntersection(reward = {\n",
    "    \"forward\": forward_r,\n",
    "    \"stop\": stop_r,\n",
    "    \"collision\": collision_r,\n",
    "    \"goal\": goal_r\n",
    "})\n",
    "```\n",
    "\n",
    "Let's change the reward function and see how it affects the $\\pi^*$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "210cf7ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "tinter = TIntersection(reward={\n",
    "    \"forward\": 0.0,\n",
    "    \"stop\": -1.0,\n",
    "    \"collision\": -100,\n",
    "    \"goal\": 20\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4305d054",
   "metadata": {},
   "outputs": [],
   "source": [
    "V_star, pi_star = value_iteration(mdp=tinter, threshold=1e-5)\n",
    "print(\"V*: {}\".format(V_star))\n",
    "print(\"pi*: {}\".format(pi_star))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab503931",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_state = ... # choose an initial state\n",
    "get_solution(initial_state)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ffd29442",
   "metadata": {},
   "source": [
    "With this new reward function, it is not good almost everywhere to choose the action 0, i.e. forward."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b120a2ff",
   "metadata": {},
   "source": [
    "## 2. QMDP\n",
    "### Task 2.1: Defining new MDP\n",
    "Assume that in this new problem, we do not know exactly where the other car is. Let's introduce a new action called *look* into our problem formulation. Everytime we choose the action *look*, no car moves, we get a $-1$ reward, and receives the observation $z \\in \\{1 \\dots 5\\}$ corresponding to the position of the other car.\n",
    "\n",
    "The first thing we will do to approximate this POMDP with QMDP is to solve for the underlying MDP formulation (assuming that all states are fully observable).\n",
    "\n",
    "Tasks:\n",
    "1. Edit the class `TIntersection()` to include the new information mentioned\n",
    "2. Use value iteration/policy iteration and solve for $\\hat{V}^{\\text{MDP}}$\n",
    "\n",
    "**NOTE**: You can copy the value of all the missing probability values from the previous questions (denotes `MISSING_VALUE`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a90ac93",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TIntersectionQMDP(MDP):\n",
    "    def __init__(self, reward={\"forward\": -1, \"stop\": -5, \"collision\": -10, \"goal\": 5.0, \"look\": -1}):\n",
    "        self.states = [\n",
    "            [\"ego_{}\".format(x) for x in list(range(1, 6))],\n",
    "            [\"car_{}\".format(x) for x in list(range(1, 6))]\n",
    "        ]\n",
    "        \n",
    "        # 1. Add another action \"look\" to self.actions\n",
    "        self.actions = [\"forward\", \"stop\"]\n",
    "\n",
    "        self.gam = 0.9\n",
    "        super().__init__(\n",
    "            states=self.states, actions=self.actions, method=\"add\")\n",
    "        self.reward = reward\n",
    "        self.populate_data()\n",
    "        \n",
    "    def populate_data(self):\n",
    "        for i in range(1, 6):\n",
    "            for j in range(1, 6):\n",
    "                # add route for forward action\n",
    "                # (i, j) --> (i+1, j)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+1, 5)), \"car_{}\".format(j)], MISSING_VALUE)\n",
    "                # (i, j) --> (i+1, j+1)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+1, 5)), \"car_{}\".format(min(j+1, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i+1, j+2)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+1, 5)), \"car_{}\".format(min(j+2, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i+2, j)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+2, 5)), \"car_{}\".format(j)], MISSING_VALUE)\n",
    "                # (i, j) --> (i+2, j+1)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+2, 5)), \"car_{}\".format(min(j+1, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i+2, j+2)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", [\"ego_{}\".format(min(i+2, 5)), \"car_{}\".format(min(j+2, 5))], MISSING_VALUE)\n",
    "\n",
    "                # add route for stop action\n",
    "                # (i, j) --> (i, j+2)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", [\"ego_{}\".format(i), \"car_{}\".format(min(j+2, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i, j+1)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", [\"ego_{}\".format(i), \"car_{}\".format(min(j+1, 5))], MISSING_VALUE)\n",
    "                # (i, j) --> (i, j)\n",
    "                self.add_route([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", [\"ego_{}\".format(i), \"car_{}\".format(j)], MISSING_VALUE)\n",
    "\n",
    "                self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", self.reward[\"forward\"])\n",
    "                self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", self.reward[\"stop\"])\n",
    "\n",
    "                # add route and reward for look action\n",
    "                ## YOUR CODE HERE\n",
    "                # self.add_route(...., \"look\", ...., p=...)\n",
    "                # ...\n",
    "                # self.add_reward(..., \"look\", self.reward[\"look\"])\n",
    "                ######\n",
    "                        \n",
    "                # check for collision\n",
    "                if (i in [4, 5]) and (j in [4, 5]):\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", self.reward[\"collision\"])\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", self.reward[\"collision\"])\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"look\", self.reward[\"collision\"])\n",
    "                elif (i == 5) or (j == 5):\n",
    "                    # reaching goal\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"forward\", self.reward[\"goal\"])\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"stop\", self.reward[\"goal\"])\n",
    "                    self.add_reward([\"ego_{}\".format(i), \"car_{}\".format(j)], \"look\", self.reward[\"goal\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d7dd02f",
   "metadata": {},
   "outputs": [],
   "source": [
    "tinter_qmdp = TIntersectionQMDP()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6917f9ec",
   "metadata": {},
   "source": [
    "### Task 2.2: QMDP\n",
    "Write the QMDP function to get the **next action**, taking in consideration the belief space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "407fee3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the V_star using previously built value iteration\n",
    "V_star, pi_star = value_iteration(mdp=tinter_qmdp, threshold=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e810c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def QMDP(V_star, belief, mdp:MDP=None):\n",
    "    if mdp is None:\n",
    "        raise ValueError(\"MDP cannot be None\")\n",
    "    \n",
    "    numa, nums, R, P = mdp.get_mdp()\n",
    "\n",
    "    # compute MDP-value for state-action pairs (Q)\n",
    "    ####\n",
    "    ## YOUR CODE HERE\n",
    "    raise NotImplementedError(\"Your QMDP function is empty\")\n",
    "\n",
    "    ####\n",
    "    return None"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "488719ca",
   "metadata": {},
   "source": [
    "As we know exactly where we are all the time, and we only do not know where the other car is. Let us keep track of the probability distribution of where the other car is using an array of `b = [p_1, p_2, p_3, p_4, p_5]`, with `p_i` indicates the probability that the other car is at the $i^\\text{th}$ position, $i \\in \\{1 \\dots 5\\}$.\n",
    "\n",
    "We then provide you with the function `propagate_belief()` to incorporate our current position to give you the full array of belief space $\\in \\mathbb{R}^{25}$\n",
    "\n",
    "**Example**: Let us have our current belief `b = [0.2, 0.2, 0.2, 0.2, 0.2]` and that we are at `ego_1`, running `propagate_belief(b, \"ego_1\")` gives us:\n",
    "\n",
    "```python\n",
    "propagate_belief(b, \"ego_1\")\n",
    ">>> [0.2 0.  0.  0.  0.  0.2 0.  0.  0.  0.  0.2 0.  0.  0.  0.  0.2 0.  0.\n",
    " 0.  0.  0.2 0.  0.  0.  0. ]\n",
    "```\n",
    "\n",
    "This is because the state array is in the form `[(ego_1, car_1), (ego_1, car_2), ..., (ego_2, car_1), (ego_2, car_2), ..., (ego_5, car_5)]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00fe6326",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assume that the initial belief space is {0.2, 0.2, 0.2, 0.2, 0.2} for the 5 positions that the other car can be in\n",
    "# if *look* is chosen, we will receive the observation of where the other car is, with probability distribution \n",
    "# p(pos-1) = 0.1, p(pos)=0.8, p(pos+1)=0.1 with pos is the true position\n",
    "b = [0.2, 0.2, 0.2, 0.2, 0.2]\n",
    "\n",
    "# write a function to propagate the belief space to include our position as well\n",
    "def propagate_belief(b_car, ego_state):\n",
    "    idx = int(ego_state.replace(\"ego_\", \"\")) - 1\n",
    "    b_full = np.zeros(25)\n",
    "    for i, p in enumerate(b_car):\n",
    "        b_full[i*5 + idx] = p\n",
    "    return b_full\n",
    "\n",
    "print(propagate_belief(b, \"ego_1\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d568c7fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TEST YOUR CODE\n",
    "# we know exactly where we are, and we are confident that the other car is at position 1\n",
    "belief = propagate_belief([1.0, 0.0, 0.0, 0.0, 0.0], \"ego_1\")\n",
    "next_action = QMDP(V_star, belief, mdp=tinter_qmdp)\n",
    "print(\"Best next action: {}\".format(next_action))\n",
    "\n",
    "# we know exactly where we are, and we are confident that the other car is at position 3\n",
    "belief = propagate_belief([0.0, 0.0, 1.0, 0.0, 0.0], \"ego_3\")\n",
    "next_action = QMDP(V_star, belief, mdp=tinter_qmdp)\n",
    "print(\"Best next action: {}\".format(next_action))\n",
    "\n",
    "# we know exactly where we are, and we are not that confident where the other is, with high prop that\n",
    "# it's around position 2 or 3\n",
    "belief = propagate_belief([0.0, 0.4, 0.5, 0.1, 0.0], \"ego_3\")\n",
    "next_action = QMDP(V_star, belief, mdp=tinter_qmdp)\n",
    "print(\"Best next action: {}\".format(next_action))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c9aaa673",
   "metadata": {},
   "source": [
    "### Task 2.3: Observation and belief space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03578732",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assume that the initial belief space is {0.2, 0.2, 0.2, 0.2, 0.2} for the 5 positions that the other car can be in\n",
    "# if *look* is chosen, we will receive the observation of where the other car is, with probability distribution \n",
    "# p(pos-1) = 0.1, p(pos)=0.8, p(pos+1)=0.1 with pos is the true position\n",
    "b = [0.2, 0.2, 0.2, 0.2, 0.2]\n",
    "\n",
    "# write a function to propagate the belief space to include our position as well\n",
    "def propagate_belief(b_car, ego_state):\n",
    "    idx = int(ego_state.replace(\"ego_\", \"\")) - 1\n",
    "    b_full = np.zeros(25)\n",
    "    for i, p in enumerate(b_car):\n",
    "        b_full[i*5 + idx] = p\n",
    "    return b_full"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ccd01e13",
   "metadata": {},
   "source": [
    "We will use Bayesian inference to update our belief using the observation, and we will model our belief space as a multinomial distribution, so that we can easily do Bayesian inference with conjugate prior. In this case, it will be Dirichlet distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7238f55e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DirichletMultinominal:\n",
    "    def __init__(self, alpha=[1, 1, 1, 1, 1]):\n",
    "        self.alpha = alpha\n",
    "        self.observations = []\n",
    "    \n",
    "    def reset(self):\n",
    "        self.observations = []\n",
    "    \n",
    "    def get_posterior_predictive(self, observations=[]):\n",
    "        ### YOUR CODE HERE\n",
    "        # 1. Append the new observations to the internal observation list\n",
    "        # 2. Calculate the predictive posterior distribution and return\n",
    "        raise NotImplementedError\n",
    "        ###\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6594dbda",
   "metadata": {},
   "outputs": [],
   "source": [
    "### TEST YOUR CODE\n",
    "dist = DirichletMultinominal()\n",
    "assert dist.get_posterior_predictive() == [0.2, 0.2, 0.2, 0.2, 0.2]\n",
    "assert dist.get_posterior_predictive(observations=[0, 0, 0, 0, 0]) == [0.6, 0.1, 0.1, 0.1, 0.1]\n",
    "dist.reset()\n",
    "assert dist.get_posterior_predictive(observations=[1, 1, 1, 1, 1]) == [0.1, 0.6, 0.1, 0.1, 0.1]\n",
    "print(\"Everything is correct!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9f5c69cb",
   "metadata": {},
   "source": [
    "Run the last block to simulate our QMDP. Does the result make sense?\n",
    "\n",
    "Run it a few times, as well as changing the initial state so that you can see different results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec8b385",
   "metadata": {},
   "outputs": [],
   "source": [
    "done = False\n",
    "state = [\"ego_1\", \"car_1\"]\n",
    "state_list = [state]\n",
    "b = dist.get_posterior_predictive()\n",
    "print(\"State: {}\\tAction: {}\\tBelief: {}\".format(state, None, b))\n",
    "\n",
    "while not done:\n",
    "    belief = propagate_belief(b, state[0])\n",
    "    next_action = QMDP(V_star=V_star, belief=belief, mdp=tinter_qmdp)\n",
    "    \n",
    "    if int(next_action) == 2:\n",
    "        true_state_idx = int(state[1].replace(\"car_\", \"\")) - 1\n",
    "        obs = np.random.choice([min(true_state_idx-1, 0), true_state_idx, min(true_state_idx+1, 4)], p=[0.1, 0.8, 0.1])\n",
    "        b = dist.get_posterior_predictive([obs])\n",
    "    elif int(next_action) == 0:\n",
    "        dist.reset()\n",
    "        b = dist.get_posterior_predictive()\n",
    "\n",
    "    state_index = tinter_qmdp.get_index(state)\n",
    "\n",
    "    next_state_indices = np.where(tinter_qmdp.P[:, state_index, int(next_action)] > 0.0)[0]\n",
    "    next_state_p = tinter_qmdp.P[:, state_index, int(next_action)][next_state_indices]\n",
    "    next_state_index = np.random.choice(next_state_indices, p=next_state_p)\n",
    "    \n",
    "    state = tinter_qmdp.get_real_state_value(next_state_index)\n",
    "    state_list.append(state)\n",
    "\n",
    "    if state[0] == \"ego_5\" or state[1] == \"car_5\":\n",
    "        done = True\n",
    "    \n",
    "    print(\"State: {}\\tAction: {}\\tBelief: {}\".format(state, next_action, b))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}