exercises/templates/ex10/FunctionApproximatorControl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-71ced7bc2b82b03e",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "# Exercise 10) Function Approximators in Control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-02f5bf72a108e92e",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "After looking at function-approximator prediction we are now ready to investigate function-approximator control. This topic is still being researched and there are many interesting applications, but most of them are not yet used economically because experts in reinforcement learning are scarce. Control tasks are generally more complex than prediction tasks and we often need to apply domain-specific expert knowledge and very performant algorithms to succeed with function approximators.\n",
    "\n",
    "As in the last exercise we will have a look at the MountainCar from OpenAI's `gym`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-0f849f94f22b5741",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Please make sure to have `sklearn` installed:\n",
    "\n",
    "`pip install sklearn`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-9bd695d3e4934221",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gymnasium as gym\n",
    "import random\n",
    "from tqdm.notebook import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('seaborn-talk')\n",
    "from sklearn.kernel_approximation import RBFSampler\n",
    "import sklearn.pipeline\n",
    "import sklearn.preprocessing\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-d37cf10e4548222c",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "## 1) Artificial Neural Networks for Semi-Gradient SARSA Control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-d7889fafaa72e858",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "For the start, we will try to control the MountainCar in the same fashion that enabled us to predict the value function in the last exercise. You can refer to the last exercise to see how `PyTorch` can be used here. Write a learning algorithm that utilizes the given ANN setup to control the MountainCar. Episodes will terminate automatically after $200$ timesteps if the finish line has not been reached. Since every timestep is rewarded with $r=-1$, the goal should be to reach the finish line as fast as possible. \n",
    "\n",
    "State space: $x = \\begin{bmatrix}\\text{position} \\\\ \\text{velocity}\n",
    "\\end{bmatrix} \\in \n",
    "\\begin{bmatrix}[-1.2, 0.6] \\\\ [-0.07, 0.07]\n",
    "\\end{bmatrix}$\n",
    "\n",
    "Feature vector: $\\tilde{x} \\in \n",
    "\\begin{bmatrix}[-1, 1] \\\\ [-1, 1]\n",
    "\\end{bmatrix}$ (you may want to change or extend these features)\n",
    "\n",
    "Input space: $u \\in \\{ 0 \\text{ (left)}, 1 \\text{ (idle)}, 2 \\text{ (right)}\\}$\n",
    "\n",
    "We will give the same features as in the last exercise: simply MinMax-normalized states. Feel free to enhance this feature vector in order to improve the trainability of this algorithm.\n",
    "\n",
    "As explained, control problems with the use of function approximators tend to be quite hard. If you have the feeling of wasting your time, please proceed to the next task before giving up this exercise.\n",
    "\n",
    "![](https://media.giphy.com/media/xUPJPgq5NcxbKHrcXK/giphy.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-76d1ee608451b89f",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The following two cells import the `plot_surface` function and implement the `featurize` function. The former can be used to monitor the quality of the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import plot_surface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-4335ae0661f1e127",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def featurize(state):\n",
    "    \"\"\"Normalizes the state. Optionally also addes additional features.\"\"\"\n",
    "\n",
    "    pos = state[0]\n",
    "    pos_min = -1.2\n",
    "    pos_max = +0.6\n",
    "    \n",
    "    vel = state[1]\n",
    "    vel_min = -0.07\n",
    "    vel_max = +0.07\n",
    "    \n",
    "    norm_pos = (pos - pos_min) / (pos_max - pos_min) * 2 - 1\n",
    "    norm_vel = (vel - vel_min) / (vel_max - vel_min) * 2 - 1\n",
    "    \n",
    "    return np.array([[norm_pos, norm_vel]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model definition\n",
    "class FeedForwardNetwork(nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, output_dim):\n",
    "        super(FeedForwardNetwork, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_dim, 64)\n",
    "        self.fc2 = nn.Linear(64, 64)\n",
    "        self.fc3 = nn.Linear(64, output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Prediction of the action-values for all actions in the given state.\n",
    "        \n",
    "        This means that for an input state the output has as many elements as\n",
    "        there are actions in the environment. Each output is the predicted\n",
    "        action-value of the corresponding action in the given state.\n",
    "        \"\"\"\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = torch.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "def policy(state, model, deterministic, epsilon):\n",
    "    \"\"\"Decides on one of the actions in dependence of the current state.\n",
    "\n",
    "    Possible actions:\n",
    "        0: accelerate left\n",
    "        1: idle (do not accelerate)\n",
    "        2: accelerate right\n",
    "\n",
    "    Args:\n",
    "        state: The state vector\n",
    "        model: The action-value prediction model\n",
    "        deterministic: Whether actions are chosen deterministically or eps-greedily\n",
    "        epsilon: Probability for random action in eps-greedy\n",
    "\n",
    "    Returns:\n",
    "        action: The chosen state\n",
    "        action_values: The estimated action values for the current state\n",
    "    \"\"\"\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()\n",
    "\n",
    "    return action, action_values\n",
    "\n",
    "\n",
    "def interact(env, state, model, deterministic, epsilon):\n",
    "    \"\"\"Interact with the environment to get to the next state.\n",
    "    \n",
    "    Args:\n",
    "        env: The environment\n",
    "        state: The current state of the environment\n",
    "        model: The action-value prediction model\n",
    "        deterministic: Whether actions are chosen deterministically or eps-greedily\n",
    "        epsilon: Probability for random action in eps-greedy\n",
    "\n",
    "    Returns:\n",
    "        next_state: The state of the environment after interaction\n",
    "        reward: The reward for the current interaction\n",
    "        done: Whether the episode is finished\n",
    "        action: The chosen state\n",
    "        action_values: The estimated action values for the current state\n",
    "    \"\"\"\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()\n",
    "\n",
    "    return next_state, reward, done, action, action_values\n",
    "\n",
    "\n",
    "def learn(model, state, next_state, reward, done, action, gamma, epsilon, optimizer, loss_function):\n",
    "    \"\"\"Learns from the interaction via bootstrapping and backpropagation.\n",
    "\n",
    "    Args:\n",
    "        model: The action-value prediction model\n",
    "        state: The state of the environment before interaction\n",
    "        next_state: The state of the environment after interaction\n",
    "        reward: The reward for the interaction\n",
    "        done: Whether the episode is finished\n",
    "        action: The action chosen in this step\n",
    "        gamma: Discount factor\n",
    "        epsilon: Probability for random action in eps-greedy\n",
    "        optimizer: The optimizer instance,\n",
    "        loss_function: Callable loss function\n",
    "\n",
    "    Returns:\n",
    "        loss: The loss for the current learning step\n",
    "    \"\"\"\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()\n",
    "\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment setup\n",
    "env = gym.make('MountainCar-v0')\n",
    "\n",
    "# Hyperparameters\n",
    "alpha = 0.0001\n",
    "gamma = 1\n",
    "epsilon = 0.15\n",
    "nb_episodes = 300\n",
    "\n",
    "\n",
    "# Initialize model, optimizer, and loss function\n",
    "state, _ = env.reset()\n",
    "norm_state = featurize(state)\n",
    "n_states = len(norm_state[0])\n",
    "n_actions = 3\n",
    "model = FeedForwardNetwork(\n",
    "    input_dim=n_states,\n",
    "    output_dim=n_actions\n",
    ")\n",
    "optimizer = optim.SGD(model.parameters(), lr=alpha)\n",
    "mse = nn.MSELoss()\n",
    "\n",
    "# Training loop\n",
    "visited_states = []\n",
    "\n",
    "for j in tqdm(range(nb_episodes)):\n",
    "    k = 0\n",
    "    rewards = 0\n",
    "    state, _ = env.reset()\n",
    "\n",
    "    while True:\n",
    "        next_state, reward, done, action, action_values = interact(\n",
    "            env = env,\n",
    "            state=state,\n",
    "            model=model,\n",
    "            deterministic=False,\n",
    "            epsilon=epsilon\n",
    "        )\n",
    "\n",
    "        _ = learn(\n",
    "            model=model,\n",
    "            state=state,\n",
    "            next_state=next_state,\n",
    "            reward=reward,\n",
    "            done=done,\n",
    "            action=action,\n",
    "            gamma=gamma,\n",
    "            epsilon=epsilon,\n",
    "            optimizer=optimizer,\n",
    "            loss_function=mse\n",
    "        )\n",
    "\n",
    "        k += 1\n",
    "        visited_states.append(state.tolist())\n",
    "        rewards += reward\n",
    "        \n",
    "        state = next_state\n",
    "\n",
    "        if done:\n",
    "            if j % 100 == 0:\n",
    "                print(rewards)\n",
    "                print(f\"Episode {j}\")\n",
    "                print(f\"Length: {k}\")\n",
    "                plot_surface(model, n_states, np.array(visited_states), featurize)\n",
    "            break\n",
    "\n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-78b56082c9fd71b8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Use the following cell to test the greedy-execution scenario."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-53ef013dfafdb2d3",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "env = gym.make('MountainCar-v0', render_mode=\"human\")\n",
    "state, _ = env.reset()\n",
    "\n",
    "k = 0\n",
    "\n",
    "while True:\n",
    "    env.render()\n",
    "    k += 1\n",
    "\n",
    "    action, _ = policy(state, model, deterministic=True, epsilon=epsilon)\n",
    "\n",
    "    # Take the action in the environment\n",
    "    state, reward, terminated, truncated, _ = env.step(action)\n",
    "    done = terminated or truncated\n",
    "\n",
    "    if done:\n",
    "        break\n",
    "\n",
    "print(k)\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "cell-efd12b568890c007",
     "locked": false,
     "points": 0,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "source": [
    "## Solution 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "cell-4c58a1f8901890d4",
     "locked": false,
     "points": 0,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "source": [
    "Since extrapolation/generalization is very dominant with simple features, the (action-)value of many unseen states is based on the seen reward of a few states around the initial position. This will tear down the action-value estimation within the whole state space, which negates any optimistic parameter initialization. Therefore, inherent exploration based on action-values that are initialized to high values as we have seen it in the first exercises (e.g. in the MonteCarlo Racetrack scenario) will not be performed here and exploration will be halted rather soon."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-e9b5438fa6bc895c",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 2) Feature Engineering in Function Approximation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-89806f69ed38e223",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As you may have seen, on-policy control is so much more complicated than prediction because some simple features as in task 1 will not always do the trick. A problem that can often be seen in cases like this originates from the generalization aspect of the ANN. Usually, generalization is an advantage, but we have to explore the state space sufficiently to enable its validity. In the MountainCar scenario, the state space can not be explored with simple means because reaching a specific state is always bound to preconditions. \n",
    "\n",
    "In order to overcome this generalization dilemma, a very different feature encoding can be employed. This can be done in the form of e.g. radial basis functions (RBFs), using RBF-feature \n",
    "$x_i = e^{-\\frac{||x - c_\\mathrm{i}||^2}{2 \\sigma_\\mathrm{i}^2}}$.\n",
    "\n",
    "For more information, you can read up in the Sutton and Barto (page 221 ff.).\n",
    "\n",
    "![](RBF.png)\n",
    "(Source: Reinforcement Learning, Sutton&Barto, p.221)\n",
    "\n",
    "With this encoding the feature vector will get very large, meaning that the dependency of the ANN output on single feature components will decrease. Contrary to just using the states as features, this will decrease the ANNs ability to extrapolate to unseen areas, because a very different set of features will be \"active\" within different regions of the state space. \n",
    "\n",
    "Show the learning curves (mean and standard deviation) as function of steps per episode! Therefore, evaluate multiple independent learning runs to retrieve the information required for plotting.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-c8d13bb7239fd192",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "observation_examples = np.array([env.observation_space.sample() for x in range(10000)])\n",
    "scaler = sklearn.preprocessing.StandardScaler()\n",
    "scaler.fit(observation_examples)\n",
    "\n",
    "featurizer = sklearn.pipeline.FeatureUnion([\n",
    "    (\"rbf0\", RBFSampler(gamma=5.0, n_components = 100)),\n",
    "    (\"rbf1\", RBFSampler(gamma=2.0, n_components = 100)),\n",
    "    (\"rbf2\", RBFSampler(gamma=1.0, n_components = 100)),\n",
    "    (\"rbf3\", RBFSampler(gamma=0.5, n_components = 100)),\n",
    "    ])\n",
    "featurizer.fit(scaler.transform(observation_examples))\n",
    "\n",
    "\n",
    "def featurize(state):\n",
    "    try:\n",
    "        scaled = scaler.transform([state])\n",
    "    except:\n",
    "        print(state)\n",
    "    featurized = featurizer.transform(scaled)\n",
    "    return featurized"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-ee66b9d4e379a0aa",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Function to measure the performance using greedy execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-ce27249f89419a42",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def performance_measure(model, env, measure_len=20, epsilon=0):\n",
    "    k_vec = np.zeros(measure_len)\n",
    "    \n",
    "    for i in range(measure_len):\n",
    "        k = 0\n",
    "        state, _ = env.reset()\n",
    "        \n",
    "        while True:\n",
    "            k += 1\n",
    "            action, _ = policy(state, model, deterministic=False, epsilon=epsilon)\n",
    "                    \n",
    "            state, reward, terminated, truncated, _ = env.step(action)\n",
    "            done = terminated or truncated\n",
    "\n",
    "            if done:\n",
    "                k_vec[i] = k\n",
    "                break\n",
    "                \n",
    "    return np.mean(k_vec), np.std(k_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "nb_episodes = 300\n",
    "nb_perf_meas = 5\n",
    "\n",
    "performance_matrix = np.zeros((nb_episodes, nb_perf_meas))\n",
    "best_performance = np.zeros((nb_perf_meas, 2))\n",
    "\n",
    "env = gym.make('MountainCar-v0')\n",
    "# Input dimension has changed due to different featurization\n",
    "state, _ = env.reset()\n",
    "norm_state = featurize(state)\n",
    "n_states = len(norm_state[0])\n",
    "n_actions = 3\n",
    "\n",
    "\n",
    "for m in tqdm(range(nb_perf_meas)):\n",
    "    \n",
    "    model = FeedForwardNetwork(input_dim=n_states, output_dim=n_actions)\n",
    "    optimizer = optim.SGD(model.parameters(), lr=alpha)\n",
    "    mse_loss = nn.MSELoss()\n",
    "\n",
    "    visited_states = []\n",
    "\n",
    "    for j in tqdm(range(nb_episodes)):\n",
    "        k = 0\n",
    "        rewards = 0\n",
    "\n",
    "        state, _ = env.reset()\n",
    "\n",
    "        while True:\n",
    "\n",
    "            # YOUR CODE HERE\n",
    "            raise NotImplementedError()\n",
    "\n",
    "            if done:\n",
    "                performance_matrix[j, m] = k\n",
    "                if j % 100 == 0:\n",
    "                    print(rewards)\n",
    "                    print(f\"Episode {j}\")\n",
    "                    print(f\"Length: {k}\")\n",
    "                    plot_surface(model, n_states, np.array(visited_states), featurize)\n",
    "\n",
    "                break\n",
    "\n",
    "    best_performance[m][0], _ = performance_measure(model, env, epsilon=0) \n",
    "    best_performance[m][1], _ = performance_measure(model, env, epsilon=0.15)\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-415a7a9e835874a7",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Show the performance over the episodes of the 10 runs (average +- standard deviadtion):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-4d112c470be120af",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots()\n",
    "epochs = list(range(300))\n",
    "perf_mean = (np.mean(performance_matrix, axis = 1))\n",
    "perf_std = np.std(performance_matrix, axis = 1)\n",
    "\n",
    "\n",
    "axs.plot(epochs, perf_mean)\n",
    "axs.fill_between(epochs, perf_mean-perf_std, perf_mean+perf_std, facecolor='r')\n",
    "axs.set_ylim([100,201])\n",
    "plt.ylabel('Steps per episode')\n",
    "plt.xlabel('Episode')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "cell-a1f34cf2d43e050d",
     "locked": false,
     "points": 0,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "source": [
    "The following plots show boxplots of the greedy and epsilon-greedy ($\\epsilon = 0.15$) execution of the 20 trained agents after 100 episodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-ef71f78bcc9d8847",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "plt.boxplot( best_performance)\n",
    "plt.ylabel('Steps per episode')\n",
    "plt.xlabel('Execution after learning, mean of all agents')\n",
    "plt.xticks([1, 2], ['greedy', 'eps-greedy'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-694f99cdffeafbcd",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "env.close()\n",
    "\n",
    "env = gym.make('MountainCar-v0', render_mode=\"human\")\n",
    "state, _ = env.reset()\n",
    "\n",
    "k = 0\n",
    "\n",
    "while True:\n",
    "    env.render()\n",
    "    k += 1\n",
    "\n",
    "    next_state, reward, done, action, action_values = interact(env, state, model, deterministic=True, epsilon=epsilon)\n",
    "    \n",
    "    state = next_state\n",
    "    \n",
    "    if done:\n",
    "        break\n",
    "\n",
    "print(k)\n",
    "env.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-1acedea65cc82bde",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "## 3)  Feature Engineering in Linear Approximation with Online LSPI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-62038af01bfd788d",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As we have seen, the RBF feature engineering allows for good trainability of the ANN. Maybe, a similar result can even be achieved with the use of a linear function approximator if this feature encoding is reused. \n",
    "\n",
    "As we cannot encode the action itself with the same RBF method, we still needed to apply a change to the `featurize` function, such that the estimate can take the action properly into account. Please investigate the new `featurize` function and try to understand the motivation behind it.\n",
    "\n",
    "Write an RLS-SARSA algorithm to train a linear function approximator that estimates the action value and enables on-policy control. \n",
    "Show the learning curves (mean and standard deviation) as function of steps per episode. Therefore, evaluate multiple independent learning runs to retrieve the information required for plotting.\n",
    "Compare the two algorithms! \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def performance_measure_LSPI(w_policy,measure_len=20, epsilon=0):\n",
    "    k_vec = np.zeros(measure_len)\n",
    "\n",
    "    for i in range(measure_len):\n",
    "\n",
    "        k = 0\n",
    "        state, _ = env.reset()\n",
    "\n",
    "        while True:\n",
    "            k += 1\n",
    "            _, action = policy(state, w_policy, env.action_space.n, True, epsilon)\n",
    "            state, reward, terminated, truncated, _ = env.step(action)\n",
    "            done = terminated or truncated\n",
    "\n",
    "            if done:\n",
    "                k_vec[i] = k\n",
    "                break\n",
    "\n",
    "    return np.mean(k_vec), np.std(k_vec)\n",
    "\n",
    "def plot_surface_LSPI(w, input_dim):\n",
    "    resolution = 100\n",
    "    pos_vec = np.linspace(-1.2, 0.6, resolution)\n",
    "    vel_vec = np.linspace(-0.07, 0.07, resolution)\n",
    "\n",
    "    pos_mat, vel_mat = np.meshgrid(pos_vec, vel_vec)\n",
    "    value_tensor = np.zeros([resolution, resolution])\n",
    "\n",
    "    for pos_idx, pos in enumerate(pos_vec):\n",
    "        for vel_idx, vel in enumerate(vel_vec):\n",
    "            feat_state, _ =  policy(np.array([pos, vel]), w, env.action_space.n, True, 0)\n",
    "            value_tensor[vel_idx, pos_idx] = np.transpose(feat_state) @ w\n",
    "\n",
    "    # Plot\n",
    "    fig = plt.subplots()\n",
    "    ax = plt.axes(projection='3d')\n",
    "    ax.plot_surface(pos_mat, vel_mat, -value_tensor, cmap=\"viridis\")\n",
    "    ax.set_xlabel('\\n\\nposition')\n",
    "    ax.set_ylabel('\\n\\nvelocity')\n",
    "    ax.set_zlabel(r'$-V_\\mathrm{greedy}$', labelpad=12)\n",
    "    ax.view_init(50, -135)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-c527672dc165b3ed",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "observation_examples = np.array([env.observation_space.sample() for x in range(10000)])\n",
    "scaler = sklearn.preprocessing.StandardScaler()\n",
    "scaler.fit(observation_examples)\n",
    "\n",
    "featurizer = sklearn.pipeline.FeatureUnion([\n",
    "    (\"rbf0\", RBFSampler(gamma=5.0, n_components = 100)),\n",
    "    (\"rbf1\", RBFSampler(gamma=2.0, n_components = 100)),\n",
    "    (\"rbf2\", RBFSampler(gamma=1.0, n_components = 100)),\n",
    "    (\"rbf3\", RBFSampler(gamma=0.5, n_components = 100)),\n",
    "    ])\n",
    "featurizer.fit(scaler.transform(observation_examples))\n",
    "\n",
    "\n",
    "def featurize(state, action):\n",
    "    \"\"\"Applies the RBF featurizer to the state and adds a featurized action as well.\"\"\"\n",
    "    action_vec = np.zeros([3, 1])\n",
    "    action_vec[action] = 1\n",
    "    \n",
    "    win = 0\n",
    "    if state[0] > 0.5:\n",
    "        win = 1\n",
    "    \n",
    "    try:\n",
    "        scaled = scaler.transform([state])\n",
    "    except:\n",
    "        print(state)\n",
    "    featurized = featurizer.transform(scaled)\n",
    "    featurized = np.reshape(featurized, (-1, 1)) # make column vector\n",
    "    \n",
    "    featurized = np.append(featurized, np.array([[1]]), axis = 0)\n",
    "    \n",
    "    featurized_vec = np.array([])\n",
    "    featurized_vec = np.expand_dims(featurized_vec, axis=-1)\n",
    "    for a in action_vec:\n",
    "        if a == 1:\n",
    "            featurized_vec = np.append(featurized_vec, featurized, axis = 0)\n",
    "        elif a == 0:\n",
    "            featurized_vec = np.append(featurized_vec, np.zeros([len(featurized), 1]), axis = 0)        \n",
    "    \n",
    "    return featurized_vec * (1 - win)\n",
    "\n",
    "\n",
    "def policy(state, w, n, deterministic, epsilon):\n",
    "    \"\"\"Decides on an action in dependence of the current state.\n",
    "\n",
    "    Args:\n",
    "        state: The current state\n",
    "        w: The parameters of the model\n",
    "        n: The number of possible actions\n",
    "        deterministic: Whether the action is chosen deterministically or epsilon-greedily\n",
    "        epsilon: Probability for random action in eps-greedy\n",
    "\n",
    "    Returns:\n",
    "        feat_states: The featurizes\n",
    "        action: The chosen action\n",
    "    \"\"\"\n",
    "    # used to store the featurized states for each possible state,\n",
    "    # the featurized state is to be used to compute the q-values\n",
    "    # and which action will be chosen is only known after the \n",
    "    # q-value calculation\n",
    "    feat_states = np.zeros([len(w), n, 1]) \n",
    "    q_value = np.zeros([n])\n",
    "\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()\n",
    "        \n",
    "    return feat_states[:, action], action\n",
    "\n",
    "\n",
    "def learn(w, P, feat_state, feat_next_state, feat_dims, gamma, _lambda):\n",
    "    \"\"\"Update the parameter set and P matrix with the RLS update rule with the data\n",
    "    from the last iteration.\n",
    "    \n",
    "    Args:\n",
    "        w: Weight parameters\n",
    "        P: P-matrix\n",
    "        feat_state: The featurized state before the interaction\n",
    "        feat_next_state: The featurized state after the interaction\n",
    "        feat_dims: The dimensionality of the featurized state\n",
    "        gamma: Discount factor\n",
    "        _lambda: Forgetting factor\n",
    "\n",
    "    Returns:\n",
    "        w: Updated parameters\n",
    "        P: Updated P-matrix\n",
    "    \"\"\"\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()\n",
    "    return w, P                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-805a7b20e5c19cc2",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gamma = 1\n",
    "# try lambda < 1, what do you see?\n",
    "_lambda = 1 \n",
    "nb_episodes = 100\n",
    "epsilon = 0.15\n",
    "k_w = 10\n",
    "nb_perf_meas = 5\n",
    "\n",
    "env = gym.make('MountainCar-v0')\n",
    "state, _ = env.reset()\n",
    "feat_state = featurize(state, 0)\n",
    "feat_dims = len(feat_state)\n",
    "\n",
    "performance_matrix_ = np.zeros((nb_episodes, nb_perf_meas))\n",
    "best_performance = np.zeros((nb_perf_meas,2))\n",
    "\n",
    "for m in tqdm(range(nb_perf_meas)):\n",
    "    k = 0\n",
    "    P = np.eye(feat_dims)\n",
    "    w = np.zeros(feat_dims)\n",
    "    w = np.expand_dims(w, axis=-1)\n",
    "\n",
    "    w_policy = np.copy(w)\n",
    "\n",
    "    for j in tqdm(range(nb_episodes)):\n",
    "        \n",
    "        length = 0\n",
    "\n",
    "        state, _ = env.reset()\n",
    "        feat_state, action = policy(state, w, env.action_space.n, False, epsilon)\n",
    "\n",
    "        while True:\n",
    "            \n",
    "            # YOUR CODE HERE\n",
    "            raise NotImplementedError()\n",
    "\n",
    "            if done:\n",
    "                performance_matrix_[j,m] = length\n",
    "                if j % 25 == 0:\n",
    "                    plot_surface_LSPI(w, feat_dims)\n",
    "                break\n",
    "                \n",
    "    best_performance[m][0], _ = performance_measure_LSPI(w_policy, epsilon=0)\n",
    "    best_performance[m][1], _ = performance_measure_LSPI(w_policy, epsilon=0.15)\n",
    "        \n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-c20d5a0acfbe704f",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "env.close()\n",
    "env = gym.make('MountainCar-v0', render_mode=\"human\")\n",
    "state, _ = env.reset()\n",
    "\n",
    "k = 0\n",
    "\n",
    "while True:\n",
    "    env.render()\n",
    "    k += 1\n",
    "    \n",
    "    _, action = policy(state, w_policy, env.action_space.n,True, 0)\n",
    "    next_state, reward, terminated, truncated, _ = env.step(action)\n",
    "    done = terminated or truncated\n",
    "    state = next_state\n",
    "    \n",
    "    if done:\n",
    "        break\n",
    "\n",
    "print(k)\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f07b3d6e621dab52",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Show the performance over the episodes of the 20 runs (average +- standard deviadtion):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f67a6f028da142ec",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots()\n",
    "epochs = list(range(100))\n",
    "perf_mean = (np.mean(performance_matrix_, axis = 1))\n",
    "perf_std = np.std(performance_matrix_, axis = 1)\n",
    "\n",
    "\n",
    "axs.plot(epochs, perf_mean)\n",
    "axs.fill_between(epochs, perf_mean-perf_std, perf_mean+perf_std, facecolor='r')\n",
    "axs.set_ylim([100,201])\n",
    "plt.ylabel('Steps per episode')\n",
    "plt.xlabel('Episode')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "cell-44bdf8ebb9fcd19f",
     "locked": false,
     "points": 0,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "source": [
    "The following plots show boxplots of the greedy and epsilon-greedy ($\\epsilon = 0.15$) execution of the 20 trained agents after 100 episodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-6e81fef7350a76be",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.boxplot( best_performance)\n",
    "plt.ylabel('Steps per episode')\n",
    "plt.xlabel('Execution after learning, mean of all agents')\n",
    "plt.xticks([1, 2], ['greedy', 'eps-greedy'])\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "cell-505c610e85afc61c",
     "locked": false,
     "points": 0,
     "schema_version": 3,
     "solution": true,
     "task": false
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The use of an appropriate RBF feature engineering enables even a linear agent to solve the MountainCar environment. Since the policy is changing over time, it would be useful to set the $\\lambda$-parameter lower than one. However, this can lead to numeric instability, so this degree of freedom has to be tuned with caution. With $\\lambda=1$, all observations are weighted equally, which does not accurately describe the behavior of our changing policy. Nonetheless, the resulting estimator is able to perform acceptable on-policy control. With a rising number of training episodes, even the $\\lambda=1$ training will converge to the correct action-value eventually (but we do not know how many episodes this may need)."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Create Assignment",
  "kernelspec": {
   "display_name": "RLKurs2",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}