6-Summary-assignments.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "## Reinforcement Learning 6: *Planning and learning with tabular methods*\n",
    "\n",
    "**Assignment:** Hand-in until 28/03/2022\n",
    "\n",
    "This week's assignment intends to summarize and compare the concepts and methods taught throughout the course.\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "import time"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "#### Circle world\n",
    "\n",
    "Imagine an agent navigating a circular environment, a simple circular track comprising `n_states=10`. At any point in time, the agent may choose to step `left` or `right`. All steps have reward $-1/(N-1)$ except for state 0 with reward 1. States 0 and $N$ are connected to form a circle.\n",
    "\n",
    "The example tasks for this assignment comprise a summary of all the content taught in the course:\n",
    " 1. *Week2*: MDPs\n",
    " 2. *Week3*: Dynamic Programming\n",
    " 3. *Week4*: Monte Carlo methods\n",
    " 4. *Week5*: TD-learning"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We start by defining the environments / tasks:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class FiniteMDP():\n",
    "    \"\"\"\n",
    "    Generic base class for MDPs with finite state, action and reward spaces\n",
    "    \"\"\"\n",
    "    def __init__(self, n_states, n_actions, reward, task, gamma=1.0, state_labels=None, action_labels=None):\n",
    "        \"\"\"\n",
    "        n_states: number of states [0,...,N-1]\n",
    "        n_actions: number of actions [0,...,N-1]\n",
    "        reward: reward values\n",
    "        task: episodic or continuing\n",
    "        gamma: discounting factor\n",
    "        \"\"\"\n",
    "        self.n_states = n_states\n",
    "        self.n_actions = n_actions\n",
    "        self.state_labels = state_labels or np.arange(self.n_states)\n",
    "        self.action_labels = action_labels or np.arange(self.n_actions)\n",
    "        self.n_rewards = len(reward)\n",
    "        self.reward = reward\n",
    "        assert(task == 'episodic' or task == 'continuing')\n",
    "        self.task = task\n",
    "        self.gamma = gamma\n",
    "\n",
    "    def p_initial(self, s):\n",
    "        \"\"\"\n",
    "        Probability for an initial state; assumed uniform\n",
    "        \"\"\"\n",
    "        return 1 / self.n_states\n",
    "\n",
    "    def sample_initial(self):\n",
    "        \"\"\"\n",
    "        Sample initial state at start of the episode; assumed uniform\n",
    "        \"\"\"\n",
    "        return np.random.randint(self.n_states)\n",
    "\n",
    "    def p_transition(self, s, a, s1, r):\n",
    "        \"\"\"\n",
    "        Transition density s x a => s1 x r\n",
    "        \"\"\"\n",
    "        # This only works for deterministic state transitions; otherwise override\n",
    "        return np.float((s1, r) == self.sample_transition(s, a))\n",
    "\n",
    "    def sample_transition(self, s, a):\n",
    "        \"\"\"\n",
    "        Sample new state and reward when starting in s and taking action a\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def sample_action(self, state, policy):\n",
    "        # sample action from policy for a given state\n",
    "        return np.random.choice(np.arange(self.n_actions), p=policy[state])\n",
    "\n",
    "    def sample_sequence(self, policy, T=None):\n",
    "        \"\"\"\n",
    "        Sample a finite horizon sequence from an MDP using some policy\n",
    "        If the tasks is continuing then we sample exactly T steps\n",
    "        If the task is episodic then we sample exactly one episode or reset until we sample T steps\n",
    "        \"\"\"\n",
    "        # sequence element is state, action, reward\n",
    "        seq = []\n",
    "        # randomly sample initial state NOTE: For exploring starts we would need to sample both states and actions. This is not needed for epsilon-greedy policies\n",
    "        s = self.sample_initial()\n",
    "\n",
    "        if self.task == 'continuing':\n",
    "            assert (T is not None)\n",
    "            for t in range(T):\n",
    "                a = self.sample_action(s, policy)\n",
    "                (s1, r) = self.sample_transition(s, a)\n",
    "                seq.append([s, a, self.reward[r]])\n",
    "                s = s1\n",
    "        else:\n",
    "            t = 0\n",
    "            while True:\n",
    "\n",
    "                if T is None and self.is_terminal(s):\n",
    "                    break\n",
    "                elif t == T:\n",
    "                    break\n",
    "\n",
    "                a = self.sample_action(s, policy)\n",
    "                if self.is_terminal(s):\n",
    "                    s1 = self.sample_initial()\n",
    "                    r = 0\n",
    "                else:\n",
    "                    (s1, r) = self.sample_transition(s, a)\n",
    "                seq.append([s, a, self.reward[r]])\n",
    "                s = s1\n",
    "                t = t+1\n",
    "        return seq\n",
    "\n",
    "    def random_deterministic_policy(self):\n",
    "        \"\"\"\n",
    "        Random choice of a deterministic action for each state\n",
    "        \"\"\"\n",
    "        return np.random.multinomial(1, [1.0 / self.n_actions for a in range(self.n_actions)], self.n_states).astype('float32')\n",
    "\n",
    "    def nonterminal_states(self):\n",
    "        \"\"\"\n",
    "        :return: The set S of nonterminal states\n",
    "        \"\"\"\n",
    "        return [s for s in range(self.n_states) if not self.is_terminal(s)]\n",
    "\n",
    "    def terminal_states(self):\n",
    "        \"\"\"\n",
    "        :return: The set S of terminal states\n",
    "        \"\"\"\n",
    "        return [s for s in range(self.n_states) if self.is_terminal(s)]\n",
    "\n",
    "    def is_terminal(self, s):\n",
    "        \"\"\"\n",
    "        flags if s is a terminal state\n",
    "        \"\"\"\n",
    "        return False  # must be implemented for episodic tasks\n",
    "\n",
    "    def __str__(self):\n",
    "        str = \"task: {0}\\n\".format(self.task)\n",
    "        str += \"states: {0}\\n\".format(self.state_labels)\n",
    "        if self.terminal_states():\n",
    "            str += \"terminal states: {0}\\n\".format(self.terminal_states())\n",
    "        str += \"actions: {0}\\n\".format(self.action_labels)\n",
    "        str += \"rewards: {0}\\n\".format(self.reward)\n",
    "        str += \"discounting factor: {0}\".format(self.gamma)\n",
    "        return str\n",
    "\n",
    "    def print_policy(self, policy):\n",
    "        for s in range(self.n_states):\n",
    "            a = np.random.choice(np.arange(self.n_actions), p=policy[s])\n",
    "            print('state ' + str(self.state_labels[s]) + ' => action ' + str(self.action_labels[a]))\n",
    "\n",
    "    def print_value(self, vf):\n",
    "        \"\"\"\n",
    "        :param vf: state value or action value function\n",
    "        \"\"\"\n",
    "        if vf.ndim == 1:\n",
    "            for s in range(self.n_states):\n",
    "                print('state ' + str(self.state_labels[s]) + ': ' + str(vf[s]))\n",
    "        else:\n",
    "            for s in range(self.n_states):\n",
    "                for a in range(self.n_actions):\n",
    "                    print('state ' + str(self.state_labels[s]) + ' - action ' + str(self.action_labels[a] + ': ' + str(vf[s,a])))\n",
    "\n",
    "    def q_to_v(self, q, policy):\n",
    "        v = np.zeros(self.n_states)\n",
    "        for s in self.nonterminal_states():\n",
    "            for a in range(self.n_actions):\n",
    "                v[s] += policy[s, a] * q[s, a]\n",
    "        return v\n",
    "\n",
    "    def v_to_q(self, v):\n",
    "        # Eqn. 4.6 in S&B\n",
    "        q = np.zeros([self.n_states, self.n_actions])\n",
    "        for s in self.nonterminal_states():\n",
    "            for a in range(self.n_actions):\n",
    "                q[s, a] = sum(self.p_transition(s, a, s1, r) * (self.reward[r] + self.gamma * v[s1])\n",
    "                              for s1 in range(self.n_states) for r in range(self.n_rewards))\n",
    "        return q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Example MDPs used in this assignment\n",
    "class CircleWorld(FiniteMDP):\n",
    "    \"\"\"\n",
    "    Walk on the circle. All steps have reward -1/(N-1)  except for state 0 with reward 1. States 0 and n_states are connected to form a circle\n",
    "    \"\"\"\n",
    "    def __init__(self, n_states, gamma=0.99):\n",
    "        super().__init__(n_states, 2, [1.0, round(-1.0 / (n_states - 1), 2)], 'continuing', gamma=gamma, state_labels=None,\n",
    "                         action_labels=['L', 'R'])\n",
    "\n",
    "    def sample_transition(self, s, a):\n",
    "        s1 = (s + 2 * a - 1) % self.n_states  # takes one step left or right\n",
    "        if s1 == 0:\n",
    "            r = 0\n",
    "        else:\n",
    "            r = 1\n",
    "        return s1, r\n",
    "\n",
    "    def optimal_policy(self):\n",
    "        \"\"\"\n",
    "        For this simple task, the optimal policy is just to move to state zero as quickly as possible\n",
    "        and then flip back and forth\n",
    "        \"\"\"\n",
    "        policy = np.zeros([self.n_states, self.n_actions])\n",
    "        n = int(np.round(self.n_states/2))\n",
    "        policy[:n, 0] = 1.0\n",
    "        policy[n:, 1] = 1.0\n",
    "        policy[self.terminal_states(), :] = 1.0 / self.n_actions\n",
    "        return policy\n",
    "\n",
    "\n",
    "class CircleWorldEpisodic(CircleWorld):\n",
    "    \"\"\"\n",
    "    Episodic variant of circle world\n",
    "    \"\"\"\n",
    "    def __init__(self, n_states, gamma=0.99):\n",
    "        super().__init__(n_states, gamma=gamma)\n",
    "        self.task = 'episodic'\n",
    "\n",
    "    def is_terminal(self, s):\n",
    "        return s == 0"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Instantiate the `CircleWorld`:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# We can instantiate it using the generic FiniteMDP class\n",
    "mdpc = FiniteMDP(n_states=10, n_actions=2, reward=[1.0, -1.0 /(10 - 1), 2], task='continuing', gamma=0.99, state_labels=None, action_labels=['L', 'R'])\n",
    "print(mdpc)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ...or we can use the example classes CircleWorld and CircleWorldEpisodic\n",
    "print('Circleworld:')\n",
    "mdpc = CircleWorld(n_states=10, gamma=0.99)\n",
    "print(mdpc)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 1. MDP simulation\n",
    "*(1 point)*\n",
    "\n",
    "**a**) Consider a uniform, stochastic policy (all actions have equal probability in all states).\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def uniform_stochastic_policy(n_states, n_actions):\n",
    "    \"\"\"\n",
    "    Each action has equal probability in all states\n",
    "    \"\"\"\n",
    "    return # TODO"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**b**) Simulate the MDP and plot the evolution of the return $G_{t}$ obtained under this policy, with discounting factors $\\gamma=1$ and $\\gamma=0.98$."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "mdpc = # TODO\n",
    "policy = # TODO\n",
    "seq = np.array(mdpc.sample_sequence(policy, 200)).T\n",
    "\n",
    "discounting = [mdpc.gamma**t for t in range(0, 200)]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(3, 1, figsize=(10, 6))\n",
    "ax[0].plot(seq[0],label='state')\n",
    "ax[0].set_xticks([])\n",
    "ax[0].set_ylabel('state')\n",
    "ax[1].plot(seq[1],label='action')\n",
    "ax[1].set_xticks([])\n",
    "ax[1].set_ylabel('action')\n",
    "ax[2].plot(np.cumsum(seq[2]),label='return ($\\gamma=$' + str(1.0) + ')' )\n",
    "ax[2].plot(np.cumsum(seq[2] * discounting),label='return ($\\gamma=$' + str(mdpc.gamma) + ')' )\n",
    "plt.legend()\n",
    "ax[2].set_xlabel('t')\n",
    "ax[2].set_ylabel('return')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 2. Dynamic Programming\n",
    "*(4 points)*\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**a**) Write a policy evaluation function, which computes the state-value function from a policy\n",
    "*(0.5 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def policy_evaluation(mdp, policy, V=None, theta=1e-8, max_t=None):\n",
    "    # Conditions for convergence\n",
    "    assert (mdp.task == 'episodic' or mdp.gamma < 1.0)\n",
    "    # Init V\n",
    "    if V is None:\n",
    "        V = np.zeros(mdp.n_states)\n",
    "    # Evaluate policy iteratively\n",
    "    t = 0\n",
    "    while True:\n",
    "        t += 1\n",
    "        delta = 0\n",
    "        for s in mdp.nonterminal_states():\n",
    "            val = V[s]\n",
    "            # TODO: correct the line below\n",
    "            V[s] = sum(policy[s, a] * mdp.p_transition(s, a, s1, r) * (mdp.reward[r] +  * V[s1])\n",
    "                       for a in range() for s1 in range() for r in range())\n",
    "            # ---\n",
    "            delta = max(delta, abs(val - V[s]))\n",
    "\n",
    "        if theta > delta or (max_t and t == max_t):\n",
    "            break\n",
    "\n",
    "    return V"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**b**) In this simple task, the optimal policy simply requires moving to state 0 as quickly as possible, so we can construct an explicit optimal policy which will be used as baseline for comparison:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# to understand the policy, study the `optimal_policy` method defined above\n",
    "pi_c = mdpc.optimal_policy()\n",
    "print(pi_c)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**b.1**) Evaluate the optimal policy generated above, using the `policy_evaluation` function you implemented above:\n",
    "*(0.5 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "V_c = # TODO: state-value function, output of policy_evaluation\n",
    "Q_c = mdpc.v_to_q(V_c)\n",
    "print(pi_c)\n",
    "print('\\nOptimal policy:')\n",
    "mdpc.print_policy(pi_c)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**b.2**) Repeat the steps above for the episodic version of the task, using `CircleWorldEpisodic`:\n",
    "*(0.5 points)*\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "print('\\nEpisodic circleworld:')\n",
    "mdpe = # TODO: instantiate MDP\n",
    "print(mdpe)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# compute state and action value of the optimal policy for the continous problem\n",
    "pi_e = # TODO: optimal policy / baseline\n",
    "V_e = # TODO: evaluate optimal policy\n",
    "Q_e = mdpe.v_to_q(V_e)\n",
    "print('\\nOptimal policy:')\n",
    "mdpc.print_policy(pi_e)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**c**) Policy iteration VS value iteration"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**c.1**) Correct/complement the implementation of `policy_improvement`:\n",
    "*(0.5 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def policy_improvement(mdp, policy, V):\n",
    "    \"\"\"\n",
    "    Policy improvement operates directly on the input policy\n",
    "    \"\"\"\n",
    "    policy_stable = True\n",
    "    for s in mdp.nonterminal_states():\n",
    "        # deterministic version\n",
    "        old_action = np.argmax(policy[s])\n",
    "        action_values = np.zeros(mdp.n_actions)\n",
    "        for a in range(mdp.n_actions):\n",
    "            # TODO: correct line below\n",
    "            action_values[a] = sum(mdp.p_transition(s, a, s1, r) * (mdp.reward[r] + mdp.gamma * V[s1])\n",
    "                                   for s1 in range() for r in range())\n",
    "        best_action = np.argmax(action_values)\n",
    "\n",
    "        # Update Policy\n",
    "        policy[s, :] = 0\n",
    "        policy[s, best_action] = 1\n",
    "        policy_changed = old_action != best_action\n",
    "        if policy_changed:\n",
    "            policy_stable = False\n",
    "    return policy, policy_stable"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**c.2**) Correct/complement the implementation of `policy_iteration`:\n",
    "*(0.75 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def policy_iteration(mdp, policy):\n",
    "    \"\"\"\n",
    "    returns policy and list of list of value estimates per state and wall-clock time\n",
    "    operates directly on the input policy\n",
    "    \"\"\"\n",
    "    start = time.time()\n",
    "    V = None\n",
    "    Vs = []\n",
    "    while True:\n",
    "        V = # TODO: evaluate policy\n",
    "        policy, policy_stable = # TODO: improve policy\n",
    "        Vs.append(copy.copy(V))\n",
    "        if policy_stable:\n",
    "            break\n",
    "    end = time.time()\n",
    "\n",
    "    return policy, Vs, end - start"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**c.3**) Correct/complement the implementation of `value_iteration`:\n",
    "*(0.75 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def policy_construction(mdp, V):\n",
    "    \"\"\"\n",
    "    Derive policy from value function\n",
    "    \"\"\"\n",
    "    policy = np.ones([mdp.n_states, mdp.n_actions]) / mdp.n_actions\n",
    "    for s in range(mdp.n_states):\n",
    "        action_values = np.zeros(mdp.n_actions)\n",
    "        for a in range(mdp.n_actions):\n",
    "            action_values[a] = sum(mdp.p_transition(s, a, s1, r) * (mdp.reward[r] + mdp.gamma * V[s1])\n",
    "                                   for s1 in range(mdp.n_states) for r in range(mdp.n_rewards))\n",
    "        best_action = np.argmax(action_values)\n",
    "        # Update Policy\n",
    "        policy[s, :] = 0\n",
    "        policy[s, best_action] = 1\n",
    "    return policy"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def value_iteration(mdp, theta=1e-8, max_t=None):\n",
    "    \"\"\"\n",
    "    returns policy and list of list of value estimates per state\n",
    "    \"\"\"\n",
    "    start = time.time()\n",
    "    V = np.zeros(mdp.n_states)\n",
    "\n",
    "    # Evaluate policy iteratively\n",
    "    t = 0\n",
    "    Vs = []\n",
    "    while True:\n",
    "        t += 1\n",
    "        delta = 0\n",
    "        for s in mdp.nonterminal_states():\n",
    "            val = V[s]\n",
    "            # TODO: correct line below\n",
    "            V[s] = max(np.sum([mdp.p_transition(s, a, s1, r) * (mdp.reward[r] + )\n",
    "                               for s1 in range(mdp.n_states)]) for a in range(mdp.n_actions) for r in range(mdp.n_rewards))\n",
    "            delta = # TODO\n",
    "        Vs.append(copy.copy(V))\n",
    "        if theta > delta or (max_t and t == max_t):\n",
    "            break\n",
    "    # Output deterministic policy\n",
    "    policy = policy_construction(mdp, V)\n",
    "    end = time.time()\n",
    "\n",
    "    return policy, Vs, end - start"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**d**) Compare `policy_iteration` with `value_iteration` in solving the `CircleWorld`\n",
    "*(0.5 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# value of random policy\n",
    "V0 = policy_evaluation(mdpc, mdpc.random_deterministic_policy())"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Run policy iteration\n",
    "policy0 = mdpc.random_deterministic_policy()\n",
    "policy1, vs1, time1 = # TODO policy_iteration\n",
    "vs1.append(policy_evaluation(mdpc, policy1))  # add value of final policy"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Run value iteration\n",
    "policy2, vs2, time2 = # TODO value_iteration\n",
    "vs2 = [V0] + vs2  # add initial policy at beginning"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "print('Initial policy:')\n",
    "mdpc.print_policy(mdpc.random_deterministic_policy())\n",
    "\n",
    "print('Policy after policy iteration:')\n",
    "mdpc.print_policy(policy1)\n",
    "\n",
    "print('Policy after value iteration:')\n",
    "mdpc.print_policy(policy2)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "tidx1 = np.arange(0,time1,time1/(len(vs1)+1))[:len(vs1)]\n",
    "plt.plot(tidx1, list(map(lambda x: np.mean((x - V_c)**2), vs1)), '-o', label='policy iteration')\n",
    "tidx2 = np.arange(0,time2,time2/(len(vs2)+1))[:len(vs2)]\n",
    "plt.plot(tidx2, list(map(lambda x: np.mean((x - V_c)**2), vs2)), '-x', label='value iteration')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('MSE')\n",
    "plt.title('policy iteration')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**e**) Plot the resulting behavior."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "seq = np.array(mdpc.sample_sequence(policy1, 100)).T\n",
    "discounting = [mdpc.gamma**t for t in range(0,100)]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.subplot(3,1,1)\n",
    "plt.plot(seq[0],label='state')\n",
    "# plt.xlabel('t')\n",
    "plt.xticks([])\n",
    "plt.ylabel('state')\n",
    "plt.subplot(3,1,2)\n",
    "plt.plot(seq[1],label='action')\n",
    "# plt.xlabel('t')\n",
    "plt.xticks([])\n",
    "plt.ylabel('action')\n",
    "plt.subplot(3,1,3)\n",
    "plt.plot(np.cumsum(seq[2]),label='return ($\\gamma=1$)')\n",
    "plt.plot(np.cumsum(seq[2] * discounting),label='return ($\\gamma=$' + str(mdpc.gamma) + ')' )\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('return')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 3. Monte Carlo Methods\n",
    "*(2 point)*\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**a**) Determine the state-values under the optimal policy for the CircleWorld MDP above using on-policy, first-visit MC:\n",
    "*(1 point)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def onpolicy_firstvisit_prediction_v(mdp, policy, num_simulations=30):\n",
    "    # Conditions for convergence\n",
    "    assert(mdp.task == 'episodic')\n",
    "    V = np.zeros(mdp.n_states)\n",
    "    returns = dict()\n",
    "    Vs = []\n",
    "    for i in range(num_simulations):\n",
    "        touched = np.zeros(mdp.n_states)\n",
    "        seq = mdp.sample_sequence(policy)\n",
    "        # Traverse sequence backwards\n",
    "        G = 0\n",
    "        for state, action, reward in reversed(seq):\n",
    "            G = # TODO\n",
    "            if touched[state]:\n",
    "                returns[state][-1] = G\n",
    "            else:\n",
    "                if state not in returns:\n",
    "                    returns[state] = []\n",
    "                returns[state].append(G)\n",
    "            # Update V based on new return\n",
    "            V[state] = np.sum(returns[state]) / float(len(returns[state]))\n",
    "        Vs.append(copy.copy(V))\n",
    "    return V, Vs"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "V, Vs = onpolicy_firstvisit_prediction_v(mdpe, mdpe.optimal_policy(), num_simulations=200)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.plot(list(map(lambda x: np.mean((x - V_e)**2), Vs)), '-o', label='state value')\n",
    "plt.plot(list(map(lambda x: np.mean((mdpe.v_to_q(x) - Q_e)**2), Vs)), '-x', label='action value')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('MSE')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**b**) Starting from a uniform random policy, solve the control problem (optimize the policy) using on-policy, first-visit MC control:\n",
    "*(1 point)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def onpolicy_firstvisit_control(mdp, policy, num_simulations=30, epsilon=0.01):\n",
    "    \"\"\"\n",
    "    Run the Monte Carlo First Visit On-Policy algorithm and return the estimated\n",
    "    policy, Q (state action) values, and returns (rewards) dict.\n",
    "    Uses epsilon-soft policies instead of exploring states (the latter starts in all possible state-action pairs)\n",
    "    p. 109 of S&B\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    mdp:\n",
    "    policy: any epsilon soft policy; e.g. mdp.uniform_stochastic_policy()\n",
    "    num_simulations : int\n",
    "        Number of episodes for the policy iteration process\n",
    "    epsilon-soft\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    numpy.ndarray\n",
    "        Estimated Policy\n",
    "    numpy.ndarray\n",
    "        Estimated Q (state-action) values\n",
    "    dict\n",
    "        Rewards obtained for every state\n",
    "    \"\"\"\n",
    "    # Conditions for convergence\n",
    "    assert(mdp.task == 'episodic')\n",
    "\n",
    "    Q = np.zeros([mdp.n_states, mdp.n_actions])\n",
    "    returns = dict()\n",
    "    Qs = []\n",
    "    for i in range(num_simulations):\n",
    "        touched = np.zeros([mdp.n_states, mdp.n_actions])\n",
    "        seq = mdp.sample_sequence(policy)\n",
    "        # Traverse sequence backwards\n",
    "        G = 0\n",
    "        for state, action, reward in reversed(seq):\n",
    "            tupl = (state, action)\n",
    "            G = # TODO\n",
    "            if touched[state, action]:\n",
    "                returns[tupl][-1] = G\n",
    "            else:\n",
    "                if tupl not in returns:\n",
    "                    returns[tupl] = []\n",
    "                returns[tupl].append(G)\n",
    "            # Update Q based on new return\n",
    "            Q[state, action] = np.sum(returns[tupl]) / float(len(returns[tupl]))\n",
    "        # update epsilon soft policy\n",
    "        for state, _, _ in seq:\n",
    "            # randomly break ties\n",
    "            a_max = np.random.choice(np.flatnonzero(Q[state] == np.max(Q[state])))\n",
    "            policy[state,:] = # TODO\n",
    "            policy[state, a_max] = # TODO\n",
    "        Qs.append(copy.copy(Q))\n",
    "    return policy, Qs"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "policy0 = uniform_stochastic_policy(mdpe.n_states, mdpc.n_actions)\n",
    "policy, Qs = onpolicy_firstvisit_control(mdpe, policy0, num_simulations=10000, epsilon=0.01)\n",
    "\n",
    "print('Policy after on-policy MCMC:')\n",
    "print(policy)\n",
    "\n",
    "mdpe.print_value(Qs[-1])\n",
    "mdpe.print_value(Q_e)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.plot(list(map(lambda x: np.mean((mdpe.q_to_v(x, policy) - V_e)**2), Qs)), '-o', label='state value')\n",
    "plt.plot(list(map(lambda x: np.mean((x - Q_e)**2), Qs)), '-x', label='action value')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('MSE')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 4. TD-methods\n",
    "*(3 points)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**a**) Implement and run TD(0) to estimate the value function under the optimal policy\n",
    "*(1 point)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def td0_prediction(mdp, policy, num_simulations=30, alpha = 0.01):\n",
    "    # Conditions for convergence; note that we can also run on continuing problems\n",
    "    assert(mdp.task == 'episodic')\n",
    "    V = np.zeros(mdp.n_states)\n",
    "    Vs = []\n",
    "    for i in range(num_simulations):\n",
    "        s = mdp.sample_initial()\n",
    "        while not mdp.is_terminal(s):\n",
    "            a = mdp.sample_action(s, policy)\n",
    "            (s1, r) = mdp.sample_transition(s, a)\n",
    "            V[s] += alpha * # TODO\n",
    "            s = s1\n",
    "        Vs.append(copy.copy(V))\n",
    "    return V, Vs"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.plot(list(map(lambda x: np.mean((x - V_e)**2), Vs)), '-o', label='state value')\n",
    "plt.plot(list(map(lambda x: np.mean((mdpe.v_to_q(x) - Q_e)**2), Vs)), '-x', label='action value')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('MSE')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**b**) Implement and run the SARSA algorithm to estimate the optimal policy\n",
    "*(1 point)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def sarsa(mdp, num_simulations=30, alpha=0.01, epsilon=0.1):\n",
    "    # Conditions for convergence; note that we can also run on continuing problems\n",
    "    assert(mdp.task == 'episodic')\n",
    "    Q = np.zeros([mdp.n_states, mdp.n_actions])\n",
    "    policy = uniform_stochastic_policy(mdp.n_states, mdp.n_actions)\n",
    "    Qs = []\n",
    "\n",
    "    for t in range(num_simulations):\n",
    "        s = mdp.sample_initial()\n",
    "        a = mdp.sample_action(s, policy)\n",
    "        while not mdp.is_terminal(s):\n",
    "            (s1, r) = mdp.sample_transition(s, a)\n",
    "            a1 = mdp.sample_action(s1, policy)\n",
    "            Q[s, a] += alpha * # TODO\n",
    "\n",
    "            # update policy\n",
    "            a_max = np.random.choice(np.flatnonzero(Q[s] == np.max(Q[s])))\n",
    "            policy[s, :] = epsilon / mdp.n_actions\n",
    "            policy[s, a_max] = # TODO\n",
    "            s = s1\n",
    "            a = a1\n",
    "        Qs.append(copy.copy(Q))\n",
    "    return policy, Q, Qs"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "policy, Q, Qs = sarsa(mdpe, num_simulations=10000, alpha=0.01, epsilon=0.1)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.plot(list(map(lambda x: np.mean((mdpe.q_to_v(x, policy) - V_e)**2), Qs)), '-o', label='state value')\n",
    "plt.plot(list(map(lambda x: np.mean((x - Q_e)**2), Qs)), '-x', label='action value')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('MSE')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**c**) Implement and run the Q-learning algorithm to estimate the optimal policy.\n",
    "*(1 point)*"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def qlearning(mdp, behavioral_policy, num_simulations=30, alpha=0.01, epsilon=0.1):\n",
    "    # Conditions for convergence; note that we can also run on continuing problems\n",
    "    assert(mdp.task == 'episodic')\n",
    "\n",
    "    Q = np.zeros([mdp.n_states, mdp.n_actions])\n",
    "    Qs = []\n",
    "\n",
    "    for t in range(num_simulations):\n",
    "        s = mdp.sample_initial()\n",
    "        while not mdp.is_terminal(s):\n",
    "            a = mdp.sample_action(s, behavioral_policy)\n",
    "            (s1, r) = mdp.sample_transition(s, a)\n",
    "            Q[s, a] += alpha * # TODO\n",
    "            s = s1\n",
    "        Qs.append(copy.copy(Q))\n",
    "    # determine policy from Q function\n",
    "    target_policy = np.zeros([mdp.n_states, mdp.n_actions])\n",
    "    for state in mdp.terminal_states():\n",
    "        target_policy[state,:] = 1.0 / mdp.n_actions\n",
    "    for state in mdp.nonterminal_states():\n",
    "        a_max = np.random.choice(np.flatnonzero(Q[state] == np.max(Q[state])))\n",
    "        target_policy[state, a_max] = 1.0\n",
    "\n",
    "    return target_policy, Q, Qs"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "b_policy = uniform_stochastic_policy(mdpe.n_states, mdpe.n_actions)\n",
    "policy, Q, Qs = qlearning(mdpe, behavioral_policy=b_policy, num_simulations=2000, alpha=0.01)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.plot(list(map(lambda x: np.mean((mdpe.q_to_v(x, policy) - V_e)**2), Qs)), '-o', label='state value')\n",
    "plt.plot(list(map(lambda x: np.mean((x - Q_e)**2), Qs)), '-x', label='action value')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('MSE')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}