{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Challenge: Temporal Difference Learning"
      ],
      "metadata": {
        "id": "qF30h0MvaR36"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Introduction"
      ],
      "metadata": {
        "id": "_3erWPNLaXtt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this challenge, your goal is to implement two TD control methods: SARSA and Q-learning.\n",
        "\n",
        "You will:\n",
        "- Write SARSA and Q-learning;\n",
        "- Test your solutions in a custom environment;\n",
        "- Visually analyze action 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": "kTHTzSr7aYbP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Dependencies"
      ],
      "metadata": {
        "id": "EDN1ThFlxqTO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make sure to install and import all dependencies before you begin to work with the code"
      ],
      "metadata": {
        "id": "pyYsXycdxsmJ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KiKSCjjEUkY9"
      },
      "outputs": [],
      "source": [
        "!pip install numpy gymnasium git+https://github.com/Codefinity-ML/codefinity-rl.git"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import gymnasium as gym\n",
        "from codefinityrl.challenges.utils import VideoRecord, plot_values, plot_policy, np_random_generator\n",
        "from codefinityrl.challenges.td import *"
      ],
      "metadata": {
        "id": "-JbVEepuUlZB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Environment"
      ],
      "metadata": {
        "id": "u2fQoJFjxmgC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For this challenge, you will again use a custom environment **KeyAndChest**. Here is a small reminder of how it works.\n",
        "\n",
        "KeyAndChest 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",
        "**NOTE: this environment still has additional features, but you shouldn't use them in your solutions.**"
      ],
      "metadata": {
        "id": "Z1Bg02VrxlcW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## SARSA"
      ],
      "metadata": {
        "id": "0u9a_K4CyOwp"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement SARSA.\n",
        "\n",
        "The pseudocode for SARSA looked like this:\n",
        "<img src=\"https://codefinity-content-media-v2.s3.eu-west-1.amazonaws.com/courses/5f9425ad-7012-42d1-9fc0-d5998b1f5fd4/section-5/sarsa.png\" />\n",
        "\n",
        "Notes on implementation:\n",
        "- Action values and policy for state $s$ stay uninitialized, until the state is first encountered. `init_state` method should be used to initialize the state;\n",
        "- Intial values for any state-action pair should be 0, even for terminal states;\n",
        "- Initial action for any state should be 0, even for terminal states;\n",
        "- `update_policy` method should be used to update the policy for a state;\n",
        "- While two actions can have the same value, you should always select the action with lesser id when updating the policy;\n",
        "- Use `env.action_space.n` to get the number of available actions;\n",
        "- Use `self.np_random` if you need to generate random numbers."
      ],
      "metadata": {
        "id": "oeP3ZaHXyQaC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SARSA:\n",
        "    def __init__(self, env: gym.Env):\n",
        "        self.np_random = np_random_generator(True)\n",
        "\n",
        "        # The environment is saved as the agent's attribute\n",
        "        self.env = env\n",
        "\n",
        "        # Initialize agent's policy and values\n",
        "        self.policy = None\n",
        "        self.values = None\n",
        "\n",
        "    def init_state(self, state):\n",
        "        # Initialize state in policy and value dictionaries if unseen\n",
        "        if None:\n",
        "            self.policy[state] = None\n",
        "            for action in None:\n",
        "                self.values[(state, action)] = None\n",
        "\n",
        "    def update_policy(self, state):\n",
        "        # Improve policy by choosing the action with the highest estimated value\n",
        "        best_action = None\n",
        "        best_action_value = None\n",
        "        for action in None:\n",
        "            if None:\n",
        "                best_action_value = None\n",
        "                best_action = None\n",
        "        self.policy[state] = None\n",
        "\n",
        "    def get_action(self, state, epsilon: float = 0):\n",
        "        # Choose an action using epsilon-greedy strategy\n",
        "        if None:\n",
        "            return None # Random action\n",
        "        else:\n",
        "            return None # Current optimal action\n",
        "\n",
        "    def train(self, gamma: float, episodes: int, epsilon: float, alpha: float):\n",
        "        # Iterate over a given number of episodes\n",
        "        for _ in None:\n",
        "            state, _ = None # Get first state\n",
        "            pass # Initialize first state\n",
        "\n",
        "            action = None # Get first action\n",
        "\n",
        "            done = None # Initialize termination check\n",
        "\n",
        "            # Do updates until the episode is completed\n",
        "            while None:\n",
        "                next_state, reward, terminated, truncated, _ = None # Perform a step\n",
        "                pass # Initialize next state\n",
        "                next_action = None # Choose next action\n",
        "                self.values[(state, action)] += None # Update action value\n",
        "\n",
        "                pass # Update policy\n",
        "\n",
        "                # Move to next state and action\n",
        "                state = None\n",
        "                action = None\n",
        "                done = None # Update termination check\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, or a random action if policy for\n",
        "        # the state is not defined\n",
        "        if state not in self.policy:\n",
        "            return self.np_random.integers(self.env.action_space.n)\n",
        "        return self.get_action(state)"
      ],
      "metadata": {
        "id": "ECKmvw4LUoKn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution1()"
      ],
      "metadata": {
        "id": "d2RrQKhj3Yle"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check1(SARSA)"
      ],
      "metadata": {
        "id": "Gz4EWE2F3a4F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run the cell below to train your agent."
      ],
      "metadata": {
        "id": "CXY3pEo23jT6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "sarsa = SARSA(env)\n",
        "\n",
        "episodes = 10000\n",
        "sarsa.train(0.99, episodes, 0.3, 0.1)"
      ],
      "metadata": {
        "id": "1Kz2J6lsYUg-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "huIqnQv13mQ_"
      }
    },
    {
      "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(sarsa(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "0MDL7lM0tUUb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "-gJ5srut3rU9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(sarsa.values)\n",
        "plot_policy(sarsa.policy)"
      ],
      "metadata": {
        "id": "qED2UAIOYdL9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Even with a constant $\\varepsilon = 0.3$, SARSA was able to quickly discover the optimal path. In contrast, on-policy Monte Carlo control required a decaying $\\varepsilon$ just to begin converging toward the optimal behavior.\n",
        "\n",
        "Moreover, the estimated action values near the optimal path are significantly more accurate under SARSA for the given $\\varepsilon$-greedy policy. This improvement is a direct consequence of bootstrapping, which allows SARSA to refine its estimates using both immediate rewards and subsequent value predictions.\n",
        "\n",
        "However, the accuracy of action values remains limited for state-action pairs that lie far from the optimal path. These pairs are rarely visited, leading to less reliable estimates in those regions. This issue can only be mitigated by increasing the number of episodes, as adjusting $\\varepsilon$ would result in learning action value estimates for a different policy — corresponding to the updated value of $\\varepsilon$ — rather than improving the accuracy of the current one."
      ],
      "metadata": {
        "id": "jugcHpVe5RX3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Q-Learning"
      ],
      "metadata": {
        "id": "8VdlHIJGyR4K"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement Q-learning.\n",
        "\n",
        "The pseudocode for Q-learning looked like this:\n",
        "<img src=\"https://codefinity-content-media-v2.s3.eu-west-1.amazonaws.com/courses/5f9425ad-7012-42d1-9fc0-d5998b1f5fd4/section-5/q-learning.png\" />\n",
        "\n",
        "Notes on implementation:\n",
        "- Action values and policy for state $s$ stay uninitialized, until the state is first encountered. `init_state` method should be used to initialize the state;\n",
        "- Intial values for any state-action pair should be 0, even for terminal states;\n",
        "- Initial action for any state should be 0, even for terminal states;\n",
        "- `update_policy` method should be used to update the policy for a state;\n",
        "- While two actions can have the same value, you should always select the action with lesser id when updating the policy;\n",
        "- Use `env.action_space.n` to get the number of available actions;\n",
        "- Use `self.np_random` if you need to generate random numbers."
      ],
      "metadata": {
        "id": "eu5fkKeBzL3S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class QLearning:\n",
        "    def __init__(self, env: gym.Env):\n",
        "        self.np_random = np_random_generator(True)\n",
        "\n",
        "        # The environment is saved as the agent's attribute\n",
        "        self.env = env\n",
        "\n",
        "        # Initialize agent's policy and values\n",
        "        self.policy = None\n",
        "        self.values = None\n",
        "\n",
        "    def init_state(self, state):\n",
        "        # Initialize state in policy and value dictionaries if unseen\n",
        "        if None:\n",
        "            self.policy[state] = None\n",
        "            for action in None:\n",
        "                self.values[(state, action)] = None\n",
        "\n",
        "    def update_policy(self, state):\n",
        "        # Improve policy by choosing the action with the highest estimated value\n",
        "        best_action = None\n",
        "        best_action_value = None\n",
        "        for action in None:\n",
        "            if None:\n",
        "                best_action_value = None\n",
        "                best_action = None\n",
        "        self.policy[state] = None\n",
        "\n",
        "    def get_action(self, state, epsilon: float = 0):\n",
        "        # Choose an action using epsilon-greedy strategy\n",
        "        if None:\n",
        "            return None # Random action\n",
        "        else:\n",
        "            return None # Current optimal action\n",
        "\n",
        "    def train(self, gamma: float, episodes: int, epsilon: float, alpha: float):\n",
        "        # Iterate over a given number of episodes\n",
        "        for _ in None:\n",
        "            state, _ = None # Get first state\n",
        "            pass # Initialize first state\n",
        "\n",
        "            action = None # Get first action\n",
        "\n",
        "            done = None # Initialize termination check\n",
        "\n",
        "            # Do updates until the episode is completed\n",
        "            while None:\n",
        "                next_state, reward, terminated, truncated, _ = None # Perform a step\n",
        "                pass # Initialize next state\n",
        "                next_action = None # Choose next action\n",
        "                self.values[(state, action)] += None # Update action value\n",
        "\n",
        "                pass # Update policy\n",
        "\n",
        "                # Move to next state and action\n",
        "                state = None\n",
        "                action = None\n",
        "                done = None # Update termination check\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, or a random action if policy for\n",
        "        # the state is not defined\n",
        "        if state not in self.policy:\n",
        "            return self.np_random.integers(self.env.action_space.n)\n",
        "        return self.get_action(state)"
      ],
      "metadata": {
        "id": "qOaqkMbfYpvJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution2()"
      ],
      "metadata": {
        "id": "czROmAbU3d5h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check2(QLearning)"
      ],
      "metadata": {
        "id": "ssfN8zLn3e8Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run the cell below to train your agent."
      ],
      "metadata": {
        "id": "amjiKtz730nM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "q_learning = QLearning(env)\n",
        "\n",
        "episodes = 10000\n",
        "q_learning.train(0.99, episodes, 0.3, 0.1)"
      ],
      "metadata": {
        "id": "PVEO415Vv2cO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "0S259Eq433Mk"
      }
    },
    {
      "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(q_learning(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "-lMtKXbawBIy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "t1T_Qjvm37lD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(q_learning.values)\n",
        "plot_policy(q_learning.policy)"
      ],
      "metadata": {
        "id": "9zISfF1DwM5g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Similar to SARSA, Q-learning does not require a decaying $\\varepsilon$ to converge toward optimal behavior, unlike its Monte Carlo counterpart.\n",
        "\n",
        "Notably, even though actions are selected using an $\\varepsilon$-greedy policy, Q-learning updates are based on the greedy action, allowing it to estimate values for the greedy policy directly, without the need for importance sampling.\n",
        "\n",
        "However, the familiar issue remains: action value estimates are still inaccurate for state-action pairs that are rarely visited, particularly those far from the optimal path. Unlike SARSA, though, Q-learning offers greater flexibility — increasing $\\varepsilon$ to boost exploration does not interfere with the learning objective. This means that high $\\varepsilon$ values can be used to promote exploration."
      ],
      "metadata": {
        "id": "L_dMx_tqZEAB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to train an agent with $\\varepsilon = 1$. In this configuration, the agent selects actions completely at random, which significantly increases both the average episode length and the overall computation time. Training will take a few minutes to complete."
      ],
      "metadata": {
        "id": "X4mIVavyevvv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "q_learning_exploration = QLearning(env)\n",
        "\n",
        "episodes = 10000\n",
        "q_learning_exploration.train(0.99, episodes, 1, 0.1)"
      ],
      "metadata": {
        "id": "LDYlqNw_ei18"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "wzSJU1TnkHie"
      }
    },
    {
      "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(q_learning_exploration(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "4w6DxMLCkLXB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And finally, run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "lE4I3bNxfE1O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(q_learning_exploration.values)\n",
        "plot_policy(q_learning_exploration.policy)"
      ],
      "metadata": {
        "id": "RTn-rtBZfEbi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This time, every state-action pair has an accurate action value, and the resulting policy is optimal across all states."
      ],
      "metadata": {
        "id": "-rNiHrPzhgtz"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Congratulations! You've completed the final challenge of this course. With the knowledge you've obtained, you're ready to dive deeper into the fascinating world of reinforcement learning!"
      ],
      "metadata": {
        "id": "C4zC3IVYktjg"
      }
    }
  ]
}