{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Challenge: Dynamic Programming"
      ],
      "metadata": {
        "id": "eUAQ5kU6GXD3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Introduction"
      ],
      "metadata": {
        "id": "OxOylSb4Giqj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this challenge, your goal is to implement two DP methods: policy iteration and value iteration.\n",
        "\n",
        "You will:\n",
        "- Write policy and value iteration algorithms;\n",
        "- Test your solutions in a custom environment;\n",
        "- Visually analyze state values and policy.\n",
        "\n",
        "**NOTE: to complete the tasks, replace `None` and `pass` parts of the code with appropriate values.**\n",
        "\n",
        "**IMPORTANT: when it comes to generation of random numbers, follow the instructions in the tasks.**"
      ],
      "metadata": {
        "id": "93QPU4Ph0BJB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Dependencies"
      ],
      "metadata": {
        "id": "lrrJid91GnF5"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-9R4hB9dr0c"
      },
      "source": [
        "Make sure to install and import all dependencies before you begin to work with the code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "TdTNUq78tfHL"
      },
      "outputs": [],
      "source": [
        "!pip install numpy gymnasium git+https://github.com/Codefinity-ML/codefinity-rl.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NVwL1wwekv_c"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import gymnasium as gym\n",
        "from codefinityrl.challenges.utils import VideoRecord, plot_values, plot_policy\n",
        "from codefinityrl.challenges.dp import *"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Environment"
      ],
      "metadata": {
        "id": "ZMiGtS4PbZXz"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For this challenge, you will use a custom environment **KeyAndChest**. This environment represents a 2D maze with key and chest. To win, the agent has to find a key and use it to open a chest. Without a key, the chest can't be opened.\n",
        "\n",
        "The state space of this environment is represented by a tuple `(x, y, key)`, where\n",
        "- `x`: x coordinate, integer\n",
        "- `y`: y coordinate, integer\n",
        "- `key`: true if agent has key, otherwise - false, boolean\n",
        "\n",
        "The agent can take one of **four actions**:\n",
        "- UP - 0\n",
        "- DOWN - 1\n",
        "- LEFT - 2\n",
        "- RIGHT - 3\n",
        "\n",
        "This environment has some additional features, that will be useful later:\n",
        "- `simulate_step(state, action)`: simulates environment's response after taking `action` in `state`. Returns the same values that `step()` would return. This can be used as the **model** of the environment;\n",
        "- `states()`: generator function, returns every valid non-terminal state;\n",
        "- `terminal_states()`: generator function, returns every valid terminal state;\n",
        "- `actions()`: generator function, returns every valid action.\n",
        "\n",
        "This additional features can't be accessed from environment's wrappers. To access them, you need to use the `unwrapped` property of the wrapper.\n",
        "\n",
        "For example, if environment is created like this: `env = gym.make(...)`, `states()` can be accessed like this: `env.unwrapped.states()`."
      ],
      "metadata": {
        "id": "3uAdazFq5M5N"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Environment rendering"
      ],
      "metadata": {
        "id": "kubzNixV9GVc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Sometimes, you may want to look at the performance of your agent visually. Maybe you want to analyze the mistakes your agent makes, or compare different agents, or something completely different. No matter the reason, the approach remains the same:\n",
        "1. Provide the `render_mode`. Gymnasium supports different render modes: `human`, `rgb_array`, `rgb_array_list`, `ansi`. Choose the one that is both supported by the environment and is suited for your needs.\n",
        "2. Use `render()` to get the visual information(not necessary for `human` mode).\n",
        "\n",
        "Run the code cell below to see a recording of how the random agent solves the environment. Notice the `rgb_array` render mode and `VideoRecord` class. `VideoRecord` is a custom class that stores frames in a list, and displays a video on `play()`. Usually, you would use `human` render mode in this situation, but notebooks don't support it."
      ],
      "metadata": {
        "id": "ju2IqLQM5OyO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\", render_mode=\"rgb_array\")\n",
        "env.reset()\n",
        "\n",
        "video = VideoRecord(env.metadata[\"render_fps\"])\n",
        "video.add_frame(env.render())\n",
        "done = False\n",
        "while not done:\n",
        "    action = env.action_space.sample()\n",
        "    _, _, terminated, truncated, _ = env.step(action)\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "pbl0yp1fbQiK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Your solution"
      ],
      "metadata": {
        "id": "V-4xkvHROWTr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Before attempting to sove the environment with DP, try to solve it yourself, by writing a sequence of actions. Write the actions in `solution_actions` variable as a list. To check how the agent moves — run the cell.\n",
        "\n",
        "Remember, the agent can take one of **four actions**:\n",
        "- UP - 0\n",
        "- DOWN - 1\n",
        "- LEFT - 2\n",
        "- RIGHT - 3\n",
        "\n",
        "Actions should be stored in a list as numbers.\n",
        "\n",
        "**Don't forget to run the cell before submitting the solution.**"
      ],
      "metadata": {
        "id": "eWqjEekXOvGo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Write your actions in this variable\n",
        "solution_actions = None\n",
        "\n",
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\", render_mode=\"rgb_array\")\n",
        "env.reset()\n",
        "\n",
        "video = VideoRecord(env.metadata[\"render_fps\"])\n",
        "video.add_frame(env.render())\n",
        "for action in solution_actions:\n",
        "    env.step(action)\n",
        "    video.add_frame(env.render())\n",
        "video.play()"
      ],
      "metadata": {
        "id": "jtYyzIbO8Amg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution1()"
      ],
      "metadata": {
        "id": "EZvBBGzIWgyi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check1(solution_actions)"
      ],
      "metadata": {
        "id": "JpO1OKN_WiWf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Policy Iteration"
      ],
      "metadata": {
        "id": "U8dOUm_Tt4fd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement a **policy iteration**.\n",
        "\n",
        "The pseudocode for the policy iteration looked like this:\n",
        "\n",
        "<img src=\"https://codefinity-content-media-v2.s3.eu-west-1.amazonaws.com/courses/5f9425ad-7012-42d1-9fc0-d5998b1f5fd4/section-3/policy-iteration-pseudocode.png\"/>\n",
        "\n",
        "While the core idea should be understandable from the pseudocode, some things need to be clarified.\n",
        "\n",
        "Notes on implementation:\n",
        "- Intial values for any state should be 0;\n",
        "- Values for terminal states are never updated, they should always stay zero;\n",
        "- Initial action for any non-terminal state should be 0;\n",
        "- Policy shouldn't define actions for terminal states;\n",
        "- While two actions can have the same value, you should always select the action with lesser id during policy improvement."
      ],
      "metadata": {
        "id": "cXuYV3UYTO6z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PolicyIterationAgent:\n",
        "    def __init__(self, env: gym.Env):\n",
        "        # To omit writing env.unwrapped every time, additional functions are\n",
        "        # saved as agent's attributes\n",
        "        self.states = env.unwrapped.states\n",
        "        self.terminal_states = env.unwrapped.terminal_states\n",
        "        self.actions = env.unwrapped.actions\n",
        "        self.model = env.unwrapped.simulate_step\n",
        "\n",
        "        # Initialize policy and values\n",
        "        self.policy = None\n",
        "        self.values = None\n",
        "        for state in self.states():\n",
        "            self.values[state] = None\n",
        "            self.policy[state] = None\n",
        "        for state in self.terminal_states():\n",
        "            self.values[state] = None\n",
        "\n",
        "    def evaluate_policy(self, theta: float, gamma: float):\n",
        "        # Write policy evaluation step in this function\n",
        "        delta = None # Initialize this correctly\n",
        "\n",
        "        # Update values while delta is high\n",
        "        while None:\n",
        "            delta = None # Reset delta\n",
        "\n",
        "            # Update the value of each state according to the\n",
        "            # algorithm described in the pseudocode\n",
        "            for state in None:\n",
        "                v = None\n",
        "                next_state, reward, _, _, _ = None\n",
        "                self.values[state] = None\n",
        "                delta = None\n",
        "\n",
        "    def improve_policy(self, gamma: float) -> bool:\n",
        "        # Write policy improvement step in this function\n",
        "        is_policy_stable = None # Initialize this correctly\n",
        "\n",
        "        # Update the action for each state\n",
        "        for state in None:\n",
        "            a = None\n",
        "\n",
        "            # Look for the best action from all options\n",
        "            best_action = None\n",
        "            best_action_value = None\n",
        "            for action in None:\n",
        "                next_state, reward, _, _, _ = None\n",
        "                action_value = None\n",
        "                if None:\n",
        "                    best_action = None\n",
        "                    best_action_value = None\n",
        "            self.policy[state] = None\n",
        "\n",
        "            # Check if policy remained stable\n",
        "            if None:\n",
        "                is_policy_stable = None\n",
        "        return None # Return if policy is stable\n",
        "\n",
        "    def train(self, theta: float, gamma: float):\n",
        "        # Combine policy evaluation and policy improvement\n",
        "        is_policy_stable = None # Initialize this correctly\n",
        "\n",
        "        # Write a correct training loop\n",
        "        while None:\n",
        "            pass # Do policy evaluation\n",
        "            is_policy_stable = None # Do policy improvement\n",
        "\n",
        "    def __call__(self, state):\n",
        "        # Calling this agent with a state will return the action that the agent\n",
        "        # chooses according to the policy\n",
        "        return self.policy[state]"
      ],
      "metadata": {
        "id": "Cx-9jVKZt30P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution2()"
      ],
      "metadata": {
        "id": "0x1OYQOtWnhi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check2(PolicyIterationAgent)"
      ],
      "metadata": {
        "id": "RRM9O1pWWoe8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run the cell below to train your agent."
      ],
      "metadata": {
        "id": "YULRgko9AdO-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "\n",
        "pi_agent = PolicyIterationAgent(env)\n",
        "pi_agent.train(1e-6, 0.99)"
      ],
      "metadata": {
        "id": "M8mVy4BlBNga"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "KBJseRTDCBm8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\", render_mode=\"rgb_array\")\n",
        "state, _ = env.reset()\n",
        "\n",
        "video = VideoRecord(env.metadata[\"render_fps\"])\n",
        "video.add_frame(env.render())\n",
        "done = False\n",
        "while not done:\n",
        "    state, _, terminated, truncated, _ = env.step(pi_agent(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "l-seeTApKODF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "HKB_qGpaDRsm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(pi_agent.values)\n",
        "plot_policy(pi_agent.policy)"
      ],
      "metadata": {
        "id": "ejrb7Ih2hA-_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the first row, you can see the value function for each state, and a map in the end for layout reference. As you can see, each state has an accurate value, even if the state is not a part of any optimal route.\n",
        "\n",
        "In the second row, you can see how the policy looks for each state, and a map in the end for layout reference. And like with the values, following this policy from any state will result in an optimal path from that state."
      ],
      "metadata": {
        "id": "y10pT_eUG9pP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Value iteration"
      ],
      "metadata": {
        "id": "iTs6ozhY77S7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement a **value iteration**.\n",
        "\n",
        "The pseudocode for the value iteration looked like this:\n",
        "\n",
        "<img src=\"https://codefinity-content-media-v2.s3.eu-west-1.amazonaws.com/courses/5f9425ad-7012-42d1-9fc0-d5998b1f5fd4/section-3/value-iteration-pseudocode.png\"/>\n",
        "\n",
        "While the core idea should be understandable from the pseudocode, like with the policy iteration, some things need to be clarified.\n",
        "\n",
        "Notes on implementation:\n",
        "- Intial values for any state should be 0;\n",
        "- Values for terminal states are never updated, they should always stay zero;\n",
        "- Policy should be undefined until policy improvement step;\n",
        "- Policy shouldn't define actions for terminal states;\n",
        "- While two actions can have the same value, you should always select the action with lesser id during policy improvement."
      ],
      "metadata": {
        "id": "rz0O7oQsHPoF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ValueIterationAgent:\n",
        "    def __init__(self, env: gym.Env):\n",
        "        # To omit writing env.unwrapped every time, additional functions are\n",
        "        # saved as agent's attributes\n",
        "        self.states = env.unwrapped.states\n",
        "        self.terminal_states = env.unwrapped.terminal_states\n",
        "        self.actions = env.unwrapped.actions\n",
        "        self.model = env.unwrapped.simulate_step\n",
        "\n",
        "        # Initialize policy and values\n",
        "        self.policy = None\n",
        "        self.values = None\n",
        "        for state in self.states():\n",
        "            pass\n",
        "        for state in self.terminal_states():\n",
        "            pass\n",
        "\n",
        "    def evaluate_policy(self, theta: float, gamma: float):\n",
        "        # Write policy evaluation step in this function\n",
        "        delta = None # Initialize this correctly\n",
        "\n",
        "        # Update values while delta is high\n",
        "        while None:\n",
        "            delta = None # Reset delta\n",
        "\n",
        "            # Update the value of each state according to the\n",
        "            # algorithm described in the pseudocode\n",
        "            for state in None:\n",
        "                v = None\n",
        "                max_q = None\n",
        "\n",
        "                # Find a maximal action value\n",
        "                for action in self.actions():\n",
        "                    next_state, reward, _, _, _ = None\n",
        "                    max_q = None\n",
        "\n",
        "                self.values[state] = None\n",
        "                delta = None\n",
        "\n",
        "    def improve_policy(self, gamma: float) -> bool:\n",
        "        # Write policy improvement step in this function\n",
        "\n",
        "        # Update the action for each state\n",
        "        for state in None:\n",
        "            a = None\n",
        "\n",
        "            # Look for the best action from all options\n",
        "            best_action = None\n",
        "            best_action_value = None\n",
        "            for action in None:\n",
        "                next_state, reward, _, _, _ = None\n",
        "                action_value = None\n",
        "                if None:\n",
        "                    best_action = None\n",
        "                    best_action_value = None\n",
        "            self.policy[state] = None\n",
        "\n",
        "            # Check if policy remained stable\n",
        "            if None:\n",
        "                is_policy_stable = None\n",
        "        return None # Return if policy is stable\n",
        "\n",
        "    def train(self, theta: float, gamma: float):\n",
        "        # Combine policy evaluation and policy improvement\n",
        "        pass # Do policy evaluation\n",
        "        pass # Do policy improvement\n",
        "\n",
        "    def __call__(self, state):\n",
        "        # Calling this agent with a state will return the action that the agent\n",
        "        # chooses according to the policy\n",
        "        return self.policy[state]"
      ],
      "metadata": {
        "id": "Qq5u_bZg79w6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution3()"
      ],
      "metadata": {
        "id": "6gvZu_oUWtzn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check3(ValueIterationAgent)"
      ],
      "metadata": {
        "id": "JolXCnaiWu7h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run the cell below to train your agent."
      ],
      "metadata": {
        "id": "wzN1FKoSJg5J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "\n",
        "vi_agent = ValueIterationAgent(env)\n",
        "vi_agent.train(1e-6, 0.99)"
      ],
      "metadata": {
        "id": "CsDlO_UimQpP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "3fRxjZbeJ25O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\", render_mode=\"rgb_array\")\n",
        "state, _ = env.reset()\n",
        "\n",
        "video = VideoRecord(env.metadata[\"render_fps\"])\n",
        "video.add_frame(env.render())\n",
        "done = False\n",
        "while not done:\n",
        "    state, _, terminated, truncated, _ = env.step(vi_agent(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "j-W32J9ymhXx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "MKx1et4KKNCM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(vi_agent.values)\n",
        "plot_policy(vi_agent.policy)"
      ],
      "metadata": {
        "id": "RF_OGHW3mbW4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You may notice that the plots above are the same as **policy iteration** plots. This happens because both of these methods converge to the optimal value function and optimal policy. There may be many optimal policies, but the way your agent prioritizes actions with lesser id guarantees that it will always arrive at the same optimal policy."
      ],
      "metadata": {
        "id": "qmg877vqKN7w"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Congratulations! You've completed the challenge."
      ],
      "metadata": {
        "id": "Tr7CZPfeMmoe"
      }
    }
  ]
}