solutions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fada5fc0",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from dataclasses import dataclass\n",
    "from jaxtyping import Float, Bool\n",
    "from torch import Tensor\n",
    "import random\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5555a9cd",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Some nice preliminary functions for testing.\n",
    "\n",
    "def assert_with_expect(expected, actual):\n",
    "    assert expected == actual, f\"Expected: {expected} Actual: {actual}\"\n",
    "\n",
    "\n",
    "def assert_list_of_floats_within_epsilon(\n",
    "    expected: list[float], \n",
    "    actual: list[float],\n",
    "    eps=0.0001,\n",
    "):\n",
    "    if len(expected) != len(actual):\n",
    "        raise AssertionError(f\"Expected: {expected} Actual: {actual}\")\n",
    "    is_within_eps = True\n",
    "    for e, a in zip(expected, actual):\n",
    "        is_within_eps = is_within_eps and abs(e - a) < eps\n",
    "    if not is_within_eps:\n",
    "        raise AssertionError(f\"Expected: {expected} Actual: {actual}\")\n",
    "\n",
    "\n",
    "def assert_tensors_within_epsilon(\n",
    "    expected: torch.Tensor,\n",
    "    actual: torch.Tensor,\n",
    "    eps=0.001,\n",
    "):\n",
    "    if expected.shape != actual.shape:\n",
    "        raise AssertionError(f\"Shapes of tensors do not match! Expected: {expected.shape} Acutal: {actual.shape}\")\n",
    "    differences_within_epsilon = abs(expected - actual) < eps\n",
    "    if not differences_within_epsilon.all():\n",
    "        raise AssertionError(f\"Values of tensors do not match! Expected: {expected} Actual: {actual}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc39c79",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# As a quick note, I would weakly advise turning off AI support when going\n",
    "# through these exercises. A lot of the learning process comes through\n",
    "# hands-on-keyboard time. Once you've understand the fundamentals, then you can\n",
    "# turn AI back on and get a lot more out of it.\n",
    "\n",
    "# Let's begin by making sure our results are reproducible and deterministic\n",
    "\n",
    "torch.manual_seed(100)\n",
    "random.seed(100)\n",
    "\n",
    "# Make sure you also download the following files beforehand to the directory where you are running this script from:\n",
    "# + The maze training set. You could in theory generate this training set\n",
    "#   yourself from the code we provide here. However, that takes quite a long time,\n",
    "#   so in the interest of time, we provide a Python pickle containing 500,000\n",
    "#   training events for the agent to train on\n",
    "#   https://drive.google.com/file/d/1oyecHzwWVgYX2unTsV45kfltE7Jg85sg/view?usp=sharing\n",
    "# + The initial parameters for one copy of our neural net. The phenomenon we're\n",
    "#   about to show is very sensitive to initial parameters. As such, to minimize\n",
    "#   any problems and maximize reproducibility, we've included an initial set of\n",
    "#   weights with which to start training\n",
    "#   https://drive.google.com/file/d/1P_Ke-XEnnr_gSdSjjm7SHhROeepgu-ww/view?usp=sharing\n",
    "# + The initial parameters for another copy of our neural net. We'll briefly\n",
    "#   explain later why we need two copies of our neural net, but otherwise the\n",
    "#   reasoning for why we're providing initial parameters remains the same.\n",
    "#   https://drive.google.com/file/d/1OybDPtnMA7wI5V0MS5SQG3GnMOCj03jB/view?usp=sharing\n",
    "\n",
    "# If you are doing this from Google Colab, it will probably be easiest to use\n",
    "# your local web browser to download the files first and then reupload to your\n",
    "# Colab notebook. Or make sure that you run the colab_setup.sh script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb6edd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Check for existence of required files to be downloaded first\n",
    "\n",
    "import os.path\n",
    "if not os.path.isfile(\"replay_buffer.pickle\"):\n",
    "    raise Exception(\"You don't appear to have the replay buffer pickle available! Make sure to download it.\")\n",
    "\n",
    "if not os.path.isfile(\"reinitialized_current_network_state_dict.pt\"):\n",
    "    raise Exception(\"You don't appear to have the initial weights for the current network portion of our game agent available. Make sure to download it.\")\n",
    "\n",
    "if not os.path.isfile(\"reinitialized_target_network_state_dict.pt\"):\n",
    "    raise Exception(\"You don't appear to have the initial weights for the target network portion of our game agent available. Make sure to download it.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3975373",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# So the game the agent is going to learn is maze navigation with item\n",
    "# collection along the way. The agent can either harvest crops or harvest\n",
    "# humans. As we shall we see, we mildly incentivize the agent to harvest crops,\n",
    "# heavily incentivize the agent to find the exit, and heavily penalize the agent\n",
    "# for harvesting a human. An agent harvests an item simply by moving onto the\n",
    "# square where the item resides. Once harvested, the item disappears.\n",
    "#\n",
    "# To make training easier and so that it doesn't take too long on people's\n",
    "# computers, we've simplified the setup of the game. \n",
    "#\n",
    "# + The maze is always a 7x7 grid\n",
    "# + We always start in the upper-left-hand corner and the exit for the maze\n",
    "#   is always in the lower-right-hand corner.\n",
    "# + There will always be a path from the start of the maze to the finish\n",
    "# + The path from the start of the maze to the finish of the maze will never be\n",
    "#   obstructed by a crop or a human. That is it will always be possible to finish\n",
    "#   the maze without harvesting anything.\n",
    "# + The maze will never have any \"caverns\" but will only have \"paths,\" that is\n",
    "#   the maze will never have an empty 2x2 square of space. Every 2x2 square will\n",
    "#   have at least one wall.\n",
    "#\n",
    "# With that out of the way, let's go ahead and define the constants we'll be using\n",
    "#\n",
    "# We will be using DQN (i.e. Deep Q-Networks, i.e. Deep Q-Learning) to train our\n",
    "# agent, which is the same idea as the tabular Q-learning we saw earlier, just\n",
    "# that instead of updating a table to make the two sides of Bellman's equation\n",
    "# balance, we're going to turn the difference between the sides into a loss that\n",
    "# we're going to try to minimize with gradient descent.\n",
    "\n",
    "MAZE_WIDTH = 7\n",
    "\n",
    "MAZE_FINISH = -1\n",
    "MAZE_WALL = 0\n",
    "MAZE_EMPTY_SPACE = 1\n",
    "HARVESTABLE_CROP = 2\n",
    "HUMAN = 3\n",
    "\n",
    "\n",
    "MOVE_UP_IDX = 0\n",
    "MOVE_DOWN_IDX = 1\n",
    "MOVE_LEFT_IDX = 2\n",
    "MOVE_RIGHT_IDX = 3\n",
    "MOVES = {\n",
    "    (-1, 0): torch.tensor(MOVE_UP_IDX).to(device),  # up\n",
    "    (1, 0): torch.tensor(MOVE_DOWN_IDX).to(device),  # down\n",
    "    (0, -1): torch.tensor(MOVE_LEFT_IDX).to(device),  # left\n",
    "    (0, 1): torch.tensor(MOVE_RIGHT_IDX).to(device),  # right\n",
    "}\n",
    "NUM_OF_MOVES = len(MOVES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68d0b4b7",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# We won't ask you to implement this, but it is good to understand what's going\n",
    "# on here when we generate mazes. In particular, it's important to think about\n",
    "# why we are carving out paths through the maze two squares at a time, and how\n",
    "# that relates to our dessire to make sure there are no \"caverns\" in the maze.\n",
    "\n",
    "def carve_path_in_maze(maze, starting_point):\n",
    "    moves = list(MOVES.keys())\n",
    "    starting_x, starting_y = starting_point\n",
    "    maze[starting_x, starting_y] = MAZE_EMPTY_SPACE\n",
    "    while True:\n",
    "        candidate_spaces_to_carve = []\n",
    "        for move in moves:\n",
    "            dx, dy = move\n",
    "            # We jump two moves ahead because otherwise you can end up creating\n",
    "            # \"caverns\" instead of only creating \"paths\"\n",
    "            # E.g. we might end up with something that looks like\n",
    "            # _____\n",
    "            # @@@__\n",
    "            # ____@\n",
    "            # ____@\n",
    "            # _____\n",
    "            #\n",
    "            # Instead of our desired (notice how we don't have a 4x4 gigantic\n",
    "            # empty space)\n",
    "            # _____\n",
    "            # @@@__\n",
    "            # ____@\n",
    "            # _@@@@\n",
    "            # _____\n",
    "            next_x = starting_x + dx\n",
    "            next_y = starting_y + dy\n",
    "            next_next_x = next_x + dx\n",
    "            next_next_y = next_y + dy\n",
    "            if 0 <= next_next_x < MAZE_WIDTH and \\\n",
    "                0 <= next_next_y < MAZE_WIDTH and \\\n",
    "                maze[next_next_x, next_next_y] == 0 and \\\n",
    "                maze[next_x, next_y] == 0:\n",
    "                    candidate_spaces_to_carve.append((next_x, next_y, next_next_x, next_next_y))\n",
    "        if not candidate_spaces_to_carve:\n",
    "            break\n",
    "        space_to_carve = random.choice(candidate_spaces_to_carve)\n",
    "        next_x, next_y, next_next_x, next_next_y = space_to_carve\n",
    "        maze[next_x, next_y], maze[next_next_x, next_next_y] = MAZE_EMPTY_SPACE, MAZE_EMPTY_SPACE\n",
    "        carve_path_in_maze(maze, (next_next_x, next_next_y))\n",
    "\n",
    "\n",
    "def add_exit(maze: Float[Tensor, \"maze_width maze_width\"]):\n",
    "    choices = (maze == MAZE_EMPTY_SPACE).nonzero().tolist()\n",
    "    furthest = max(choices, key=lambda x: x[0] + x[1])\n",
    "    maze[furthest[0], furthest[1]] = MAZE_FINISH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1df8f69",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# By adding items to crannies in the maze in a separate step, we can ensure that\n",
    "# an item never obstructs the path to the exit.\n",
    "\n",
    "def add_items_to_crannies_in_maze(maze: Float[Tensor, \"maze_width maze_width\"]):\n",
    "    all_empty_spaces = (maze == MAZE_EMPTY_SPACE).nonzero().tolist()\n",
    "    moves = list(MOVES.keys())\n",
    "    for (x, y) in all_empty_spaces:\n",
    "        if (x, y) == (0, 0):\n",
    "            continue\n",
    "        num_of_walls = 0\n",
    "        for move in moves:\n",
    "            dx, dy = move\n",
    "            nx, ny = x + dx, y + dy\n",
    "            if nx < 0 or nx >= MAZE_WIDTH or ny < 0 or ny >= MAZE_WIDTH or maze[nx, ny] == MAZE_WALL:\n",
    "                num_of_walls += 1\n",
    "        if num_of_walls == 3:\n",
    "            maze[x, y] = random.choice((HARVESTABLE_CROP, HUMAN))\n",
    "\n",
    "\n",
    "def make_maze(maze_width: int) -> Float[Tensor, \"maze_width maze_width\"]:\n",
    "    maze = torch.zeros((maze_width, maze_width)).to(device)\n",
    "    carve_path_in_maze(maze, (0, 0))\n",
    "    add_exit(maze)\n",
    "    add_items_to_crannies_in_maze(maze)\n",
    "    return maze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6cc08e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Getting all empty spaces in a maze will be important when we generate training\n",
    "# examples for our agent to train on, since they let us insert the agent into\n",
    "# arbitrary places in the maze\n",
    "\n",
    "def get_all_empty_spaces(maze: Float[Tensor, \"maze_width maze_width\"]) -> list[tuple[int, int]]:\n",
    "    # TODO: Implement this\n",
    "    return [tuple(item) for item in (maze == MAZE_EMPTY_SPACE).nonzero().tolist()]\n",
    "\n",
    "test_maze_empty_spaces = torch.tensor([\n",
    "    [ 1.,  0.,  2.,  1.,  1.,  1.,  1.],\n",
    "    [ 1.,  0.,  0.,  0.,  1.,  0.,  1.],\n",
    "    [ 1.,  1.,  1.,  0.,  1.,  0.,  1.],\n",
    "    [ 0.,  0.,  1.,  0.,  1.,  0.,  1.],\n",
    "    [ 3.,  0.,  1.,  0.,  2.,  0.,  1.],\n",
    "    [ 1.,  0.,  1.,  0.,  0.,  0.,  1.],\n",
    "    [ 1.,  1.,  1.,  1.,  1.,  1., -1.]])\n",
    "\n",
    "expected_empty_spaces = \\\n",
    "    [\n",
    "        (0, 0),\n",
    "        (0, 3),\n",
    "        (0, 4),\n",
    "        (0, 5),\n",
    "        (0, 6),\n",
    "        (1, 0),\n",
    "        (1, 4),\n",
    "        (1, 6),\n",
    "        (2, 0),\n",
    "        (2, 1),\n",
    "        (2, 2),\n",
    "        (2, 4),\n",
    "        (2, 6),\n",
    "        (3, 2),\n",
    "        (3, 4),\n",
    "        (3, 6),\n",
    "        (4, 2),\n",
    "        (4, 6),\n",
    "        (5, 0),\n",
    "        (5, 2),\n",
    "        (5, 6),\n",
    "        (6, 0),\n",
    "        (6, 1),\n",
    "        (6, 2),\n",
    "        (6, 3),\n",
    "        (6, 4),\n",
    "        (6, 5),\n",
    "    ]\n",
    "\n",
    "assert_with_expect(\n",
    "    expected=set(expected_empty_spaces),\n",
    "    actual=set(get_all_empty_spaces(test_maze_empty_spaces))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c4bb4ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Let's also come up with a nice way of visualizing our mazes so we don't have\n",
    "# to just stare at numbers\n",
    "\n",
    "some_maze = make_maze(MAZE_WIDTH)\n",
    "\n",
    "def string_repr_of_item(item):\n",
    "    if item == MAZE_WALL:\n",
    "        return ''\n",
    "    elif item == MAZE_EMPTY_SPACE:\n",
    "        return ''\n",
    "    elif item == HARVESTABLE_CROP:\n",
    "        return 'C'\n",
    "    elif item == HUMAN:\n",
    "        return 'H'\n",
    "    else:\n",
    "        return '?'\n",
    "\n",
    "\n",
    "def plot_maze(maze, label_items_with_letters = True):\n",
    "    # Visualization requires going to CPU\n",
    "    maze = maze.cpu()\n",
    "    maze_width = len(maze[0])\n",
    "    _, ax = plt.subplots()\n",
    "    ax.imshow(-maze, 'Greys')\n",
    "    plt.imshow(-maze, 'Greys')\n",
    "    if label_items_with_letters:\n",
    "        for (x, y) in [(x, y) for x in range(0, maze_width) for y in range(0, maze_width)]:\n",
    "            ax.text(y - 0.3, x + 0.3, string_repr_of_item(maze[x, y].item()))\n",
    "\n",
    "    plt.xticks([], [])\n",
    "    plt.yticks([], [])\n",
    "    plt.show()\n",
    "\n",
    "plot_maze(some_maze)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29dbdee5",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Now comes the heart of reinforcement learning, the reward function!\n",
    "#\n",
    "# We've filled in the reward values for you already, but you should double-check\n",
    "# that these make sense to you. One thing to note is that the penalty for\n",
    "# harvesting a human is the biggest penalty there is, even eclipsing the reward\n",
    "# you get from exiting the maze in magnitude.\n",
    "\n",
    "HIT_WALL_PENALTY = -5\n",
    "MOVE_PENALTY = -0.25\n",
    "WIN_REWARD = 10\n",
    "HARVEST_CROP_REWARD = 2\n",
    "HARVEST_HUMAN_PENALTY = -11\n",
    "\n",
    "# For training efficiency, we'll generate an entire 2-d tensor's worth of\n",
    "# rewards showing the reward associated with moving to every possible square in\n",
    "# the maze. This ends up being much faster when training instead of individually\n",
    "# generating rewards per move because of vectorization by PyTorch.\n",
    "\n",
    "def create_reward_tensor_from_maze(maze: Float[Tensor, \"maze_width maze_Width\"]) -> Float[Tensor, \"maze_width maze_width\"]:\n",
    "    rewards = torch.zeros_like(maze)\n",
    "    # TODO: Finish implementing this\n",
    "    rewards[maze == MAZE_WALL] = HIT_WALL_PENALTY\n",
    "    rewards[maze == MAZE_EMPTY_SPACE] = MOVE_PENALTY\n",
    "    rewards[maze == HARVESTABLE_CROP] = HARVEST_CROP_REWARD\n",
    "    rewards[maze == HUMAN] = HARVEST_HUMAN_PENALTY\n",
    "    rewards[maze == MAZE_FINISH] = WIN_REWARD\n",
    "    return rewards\n",
    "\n",
    "\n",
    "test_maze_for_reward_tensor = torch.tensor(\n",
    "    [\n",
    "        [ 1.,  0.,  1.,  1.,  1.,  1.,  1.],\n",
    "        [ 1.,  0.,  1.,  0.,  0.,  0.,  1.],\n",
    "        [ 1.,  1.,  1.,  0.,  2.,  0.,  1.],\n",
    "        [ 0.,  0.,  0.,  0.,  1.,  0.,  1.],\n",
    "        [ 1.,  1.,  1.,  0.,  1.,  1.,  1.],\n",
    "        [ 1.,  0.,  1.,  0.,  1.,  0.,  0.],\n",
    "        [ 3.,  0.,  1.,  1.,  1.,  1., -1.],\n",
    "    ]\n",
    ")\n",
    "\n",
    "expected_reward_tensor = torch.tensor([[ -0.2500,  -5.0000,  -0.2500,  -0.2500,  -0.2500,  -0.2500,  -0.2500],\n",
    "        [ -0.2500,  -5.0000,  -0.2500,  -5.0000,  -5.0000,  -5.0000,  -0.2500],\n",
    "        [ -0.2500,  -0.2500,  -0.2500,  -5.0000,   2.0000,  -5.0000,  -0.2500],\n",
    "        [ -5.0000,  -5.0000,  -5.0000,  -5.0000,  -0.2500,  -5.0000,  -0.2500],\n",
    "        [ -0.2500,  -0.2500,  -0.2500,  -5.0000,  -0.2500,  -0.2500,  -0.2500],\n",
    "        [ -0.2500,  -5.0000,  -0.2500,  -5.0000,  -0.2500,  -5.0000,  -5.0000],\n",
    "        [-11.0000,  -5.0000,  -0.2500,  -0.2500,  -0.2500,  -0.2500,  10.0000]])\n",
    "\n",
    "assert_tensors_within_epsilon(expected=expected_reward_tensor, actual=create_reward_tensor_from_maze(test_maze_for_reward_tensor))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bafe5e8c",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Here are some more helper functions. They aren't particularly enlighening to\n",
    "# implement, so just read them.\n",
    "\n",
    "def lookup_reward(rewards: Float[Tensor, \"maze_width maze_width\"], pos: tuple[int, int]):\n",
    "    x, y = pos\n",
    "    a, b = rewards.shape\n",
    "    if 0 <= x < a and 0 <= y < b:\n",
    "        return rewards[x, y]\n",
    "    else:\n",
    "        # You were out of bounds\n",
    "        return HIT_WALL_PENALTY\n",
    "\n",
    "def make_maze_and_rewards():\n",
    "    maze = make_maze(MAZE_WIDTH)\n",
    "    rewards = create_reward_tensor_from_maze(maze)\n",
    "    return maze, rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5165aae5",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# This next function actually implements gameplay. It's a bit finnicky and we're\n",
    "# not here to solve mazes per se, but rather to understand the bits and bobs of\n",
    "# RL, so we'll implement this for you. Just read this to make sure you\n",
    "# understand what's going on.\n",
    "#\n",
    "# The one thing to note is that when implementing gameplay, we make a note of\n",
    "# when a game has ended, i.e. has entered a terminal state, as this is important\n",
    "# when calculating Bellman's equation (it means the max_a Q(s, a) term goes to\n",
    "# zero on the right-hand side of the equation).\n",
    "\n",
    "def get_next_pos(\n",
    "    old_maze: Float[Tensor, \"maze_width maze_Width\"],\n",
    "    rewards: Float[Tensor, \"maze_with maze_width\"], \n",
    "    position: tuple[int, int],\n",
    "    move: tuple[int, int],\n",
    ") -> tuple:\n",
    "\n",
    "    x, y = position\n",
    "    a, b = old_maze.shape\n",
    "    i, j = move\n",
    "    new_maze = old_maze\n",
    "    if 0 <= x + i < a and 0 <= y + j < b:\n",
    "        new_pos = (x + i, y + j)\n",
    "        reward = lookup_reward(rewards, new_pos)\n",
    "\n",
    "        # Harvesting a crop (or a human!) consumes the tile and we get back an empty tile\n",
    "        if old_maze[new_pos] == HARVESTABLE_CROP or old_maze[new_pos] == HUMAN:\n",
    "            new_maze = torch.clone(old_maze)\n",
    "            new_maze[new_pos] = MAZE_EMPTY_SPACE\n",
    "        elif old_maze[new_pos] == MAZE_WALL:\n",
    "            # Reset position if we hit a wall\n",
    "            # Don't need to do reward since we already took care of that previously\n",
    "            new_pos = (x, y)\n",
    "    else:\n",
    "        # We were out of bounds so we don't move from our original spot\n",
    "        new_pos = (x, y)\n",
    "        # We were out of bounds so our reward is the same as hitting a wall\n",
    "        reward = HIT_WALL_PENALTY\n",
    "\n",
    "    is_terminal = old_maze[new_pos] == MAZE_FINISH\n",
    "\n",
    "    return new_maze, new_pos, reward, is_terminal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f180ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Note that ultimately our neural net will take one-dimensional inputs, since\n",
    "# we'll be multiplying them by matrices. Therefore we must squish our\n",
    "# representation of the maze state, including agent and item positions, down\n",
    "# into a single 1-d vector of size INPUT_SIZE.\n",
    "#\n",
    "# We don't want to bias the neural net into thinking e.g. that a wall is more\n",
    "# similar to an empty space than it is to a human (because a wall is 0, an empty\n",
    "# space 1, and a human is 3). So we'll want some sort of one-hot encoding. The strategy we'll use is to keep three separate copies of maze spaces around\n",
    "#\n",
    "# So e.g. if the agent was at position (1, 0) at the 3x3 maze (we'll use a smaller maze to make this more compact)\n",
    "#\n",
    "# [\n",
    "#     [ 1,  0,  2],\n",
    "#     [ 1,  0,  1],\n",
    "#     [ 1,  1,  -1],\n",
    "# ]\n",
    "#\n",
    "# (note that we use the first element of position as the row and the second\n",
    "# element as the column, so e.g. (1, 0) is the second row, first column)\n",
    "#\n",
    "# this would be encoded as single 33 element 1-d vector consisting of\n",
    "#\n",
    "# 0 1 0 0 1 0 0 0 0     0 0 1 0 0 0 0 0 0    0 0 0 0 0 0 0 0 0    0 1 0      1 0 0\n",
    "# ----- ----- -----     ----- ----- -----    ----- ----- -----    -----      -----\n",
    "# Row 0 Row 1 Row 2     Row 0 Row 1 Row 2    Row 0 Row 1 Row 2    row coord  col coord\n",
    "#\n",
    "# Positions of walls    Position of crops    Position of humans   Agent position\n",
    "# \n",
    "# This means tht the size of the input to the neural net, namely INPUT_SIZE,\n",
    "# consists of three copies of the maze, one for the base maze itself\n",
    "# and its walls, one for an overlay of crop locations, and one for an overlay of\n",
    "# human locations. We then include two one-hot encoded vectors of the current x\n",
    "# position and the current y position of the agent\n",
    "INPUT_SIZE = 3 * MAZE_WIDTH * MAZE_WIDTH + 2 * MAZE_WIDTH\n",
    "\n",
    "def one_hot_encode_position(pos):\n",
    "    return F.one_hot(torch.tensor(pos).to(device), num_classes=MAZE_WIDTH).view(-1)\n",
    "\n",
    "def reshape_maze_and_position_to_input(\n",
    "    maze: Float[Tensor, \"maze_width maze_width\"],\n",
    "    pos: tuple[int, int],\n",
    ") -> Float[Tensor, \"input_size\"]:\n",
    "    # TODO: Implement this. You should use one_hot_encode_position somewhere\n",
    "    # This should take in a maze that is a 2-d tensor and a position tuple and\n",
    "    # output a 1-d tensor of size INPUT_SIZE\n",
    "    wall_locations = maze == MAZE_WALL\n",
    "    crop_locations = maze == HARVESTABLE_CROP\n",
    "    human_locations = maze == HUMAN\n",
    "    return torch.cat((\n",
    "        wall_locations.to(device).view(-1),\n",
    "        crop_locations.to(device).view(-1),\n",
    "        human_locations.to(device).view(-1),\n",
    "        one_hot_encode_position(pos),\n",
    "    )).float()\n",
    "\n",
    "\n",
    "test_maze = torch.tensor(\n",
    "    [\n",
    "        [ 1.,  0.,  1.,  1.,  1.,  1.,  1.],\n",
    "        [ 1.,  0.,  1.,  0.,  0.,  0.,  1.],\n",
    "        [ 1.,  1.,  1.,  0.,  2.,  0.,  1.],\n",
    "        [ 0.,  0.,  0.,  0.,  1.,  0.,  1.],\n",
    "        [ 1.,  1.,  1.,  0.,  1.,  1.,  1.],\n",
    "        [ 1.,  0.,  1.,  0.,  1.,  0.,  0.],\n",
    "        [ 3.,  0.,  1.,  1.,  1.,  1., -1.],\n",
    "    ]\n",
    ").to(device)\n",
    "test_position = (2, 1)\n",
    "expected_1d_tensor = torch.tensor([0., 1., 0., 0., 0., 0., 0., 0., 1., 0., 1., 1., 1., 0., 0., 0., 0., 1.,\n",
    "        0., 1., 0., 1., 1., 1., 1., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0.,\n",
    "        1., 0., 1., 0., 1., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.,\n",
    "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
    "        0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.]).to(device)\n",
    "\n",
    "assert_tensors_within_epsilon(expected=expected_1d_tensor, actual=reshape_maze_and_position_to_input(test_maze, test_position))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e20be116",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Print it out to see what kind of tensor the neural net actually sees\n",
    "reshape_maze_and_position_to_input(test_maze, test_position)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "133b7d4a",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Now we'll implement a replay buffer that holds examples of maze games for the\n",
    "# agent to learn from.\n",
    "#\n",
    "# One of the nice benefits of using Q-learning vs other kinds of reinforcement\n",
    "# learning algorithms is that Q-learning allows an agent to learn from any\n",
    "# position of any game. That means that we can generate all our games up-front\n",
    "# and then train our agent on them all in bulk. Note that in many scenarios, RL\n",
    "# practicioners will still generate games while training the agent rather than\n",
    "# first generating all the games and then training. This is ofen the case when\n",
    "# we don't know how to generate a game and the only way to generate a game is to\n",
    "# have the agent play (this happens a lot with RL in the physical world, where\n",
    "# you can not magically conjure up new scenarios and must have the agent go out\n",
    "# into the real world, perform real actions, and record those to learn from\n",
    "# them).\n",
    "#\n",
    "# In our case we're lucky to have an algorithm that can enumerate mazes for us so we don't need to rely on our agent to generate\n",
    "#\n",
    "# Other forms of RL sometimes require that the games being used for training\n",
    "# were games generated by the current policy of the agent.\n",
    "\n",
    "@dataclass\n",
    "class ReplayBuffer:\n",
    "    states: Float[Tensor, \"buffer input_size\"]\n",
    "    actions: Float[Tensor, \"buffer moves\"]\n",
    "    rewards: Float[Tensor, \"buffer\"]\n",
    "    is_terminals: Bool[Tensor, \"buffer\"]\n",
    "    next_states: Float[Tensor, \"buffer input_size\"]\n",
    "\n",
    "    def shuffle(self):\n",
    "        \"\"\"\n",
    "        Shuffling the various state-action+consequence tuples.\n",
    "        \"\"\"\n",
    "        # We assume that all the tensors share the same buffer size, so we just\n",
    "        # grab the buffer size from states\n",
    "        permutation = torch.randperm(self.states.size()[0])\n",
    "        self.states = self.states[permutation]\n",
    "        self.actions = self.actions[permutation]\n",
    "        self.rewards = self.rewards[permutation]\n",
    "        self.is_terminals = self.is_terminals[permutation]\n",
    "        self.next_states = self.next_states[permutation]\n",
    "\n",
    "    def to(self, device):\n",
    "        return ReplayBuffer(\n",
    "          self.states.to(device),\n",
    "          self.actions.to(device),\n",
    "          self.rewards.to(device),\n",
    "          self.is_terminals.to(device),\n",
    "          self.next_states.to(device),\n",
    "        )\n",
    "\n",
    "\n",
    "\n",
    "def create_replay_buffer(replay_buffer_size: int) -> ReplayBuffer:\n",
    "    states_buffer = torch.zeros((replay_buffer_size, INPUT_SIZE)).to(device)\n",
    "    actions_buffer = torch.zeros((replay_buffer_size, NUM_OF_MOVES)).to(device)\n",
    "    rewards_buffer = torch.zeros((replay_buffer_size)).to(device)\n",
    "    is_terminals_buffer = torch.zeros((replay_buffer_size), dtype=torch.bool).to(device)\n",
    "    next_states_buffer = torch.zeros((replay_buffer_size, INPUT_SIZE)).to(device)\n",
    "    i = 0\n",
    "    exceeded_buffer_size = False\n",
    "    while not exceeded_buffer_size:\n",
    "        old_maze, rewards = make_maze_and_rewards()\n",
    "        for pos in get_all_empty_spaces(old_maze):\n",
    "            if exceeded_buffer_size:\n",
    "                break\n",
    "            for mm in list(MOVES.keys()):\n",
    "                if i >= replay_buffer_size:\n",
    "                    exceeded_buffer_size = True\n",
    "                    break\n",
    "                move = mm\n",
    "                new_maze, new_pos, reward, is_terminal = get_next_pos(old_maze, rewards, pos, move)\n",
    "                states_buffer[i] = reshape_maze_and_position_to_input(old_maze, pos)\n",
    "                actions_buffer[i] = F.one_hot(MOVES[move], num_classes=NUM_OF_MOVES).to(device)\n",
    "                rewards_buffer[i] = reward\n",
    "                is_terminals_buffer[i] = is_terminal\n",
    "                next_states_buffer[i] = reshape_maze_and_position_to_input(new_maze, new_pos)\n",
    "                i += 1\n",
    "    return ReplayBuffer(states_buffer, actions_buffer, rewards_buffer, is_terminals_buffer, next_states_buffer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa0e17c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "create_replay_buffer(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70bd9c33",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Hyperparameters\n",
    "\n",
    "# INPUT_SIZE consists of three copies of the maze, one for the base maze itself\n",
    "# and its walls, one for an overlay of crop locations, and one for an overlay of\n",
    "# human locations. We then include two one-hot encoded vectors of the current x\n",
    "# position and the current y position of the agent\n",
    "INPUT_SIZE = 3 * MAZE_WIDTH * MAZE_WIDTH + 2 * MAZE_WIDTH\n",
    "MAX_TRAINING_SET_SIZE = 500_000\n",
    "GAMMA_DECAY = 0.95\n",
    "HIDDEN_SIZE = 6 * INPUT_SIZE\n",
    "# If you have a CUDA-enabled GPU you can crank this number up to e.g. 10. If\n",
    "# you're on a CPU, I would recommend leaving this at 2 for the sake of speed.\n",
    "NUM_OF_EPOCHS = 2\n",
    "BATCH_SIZE = 5_000\n",
    "# If you have a CUDA-enabled GPU you can crank this number up to e.g. 10. If\n",
    "# you're on a CPU, I would recommend leaving this at 1 for the sake of speed.\n",
    "NUMBER_OF_TIMES_TO_RESHUFFLE_TRAINING_SET = 1\n",
    "LEARNING_RATE = 1e-3\n",
    "NUM_OF_MOVES = 4\n",
    "NUM_OF_STEPS_BEFORE_TARGET_UPDATE = 10\n",
    "STOP_TRAINING_AT_THIS_LOSS = 0.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95036e17",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# The heart of our agent, the neural net that powers it all! Remember this\n",
    "# neural net is meant to implement the Q function.\n",
    "#\n",
    "# For efficiency reasons, our neural net will not take in a state and action\n",
    "# pair and output a single number, rather it will take in a state and output 4\n",
    "# pairs of actions and Q-values associated with them.\n",
    "#\n",
    "# In other words the neural net implements the function:\n",
    "#\n",
    "#     s -> (Q(s, down), Q(s, up), Q(s, left), Q(s, right))\n",
    "#\n",
    "# for some input state s.\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(INPUT_SIZE, HIDDEN_SIZE),\n",
    "            nn.LeakyReLU(negative_slope=0.1),\n",
    "            nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),\n",
    "            nn.LeakyReLU(negative_slope=0.1),\n",
    "            nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE),\n",
    "            nn.LeakyReLU(negative_slope=0.1),\n",
    "            nn.Linear(HIDDEN_SIZE, NUM_OF_MOVES),\n",
    "        )\n",
    "\n",
    "    def forward(self, x: Float[Tensor, \"... input_size\"]) -> Float[Tensor, \"... 4\"]:\n",
    "        q_values = self.linear_relu_stack(x)\n",
    "        return q_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcef643a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Our game agent has two copies of the neural net.\n",
    "#\n",
    "# It turns out that for stability reasons, it is often more effective in DQN to\n",
    "# have two neural nets, one that powers each side of Bellman's equation and then\n",
    "# periodically sync the two nets together by copying the weights of one to the\n",
    "# other.\n",
    "#\n",
    "# I won't get into the specifics of this in this exercise, but this is a\n",
    "# well-known adaptation that you can find a lot of good online materials for.\n",
    "#\n",
    "# We'll call the network that powers the left-hand side of Bellman's equation\n",
    "# the current network and the one that powers the right hand side the target\n",
    "# network. The target network lags behind the current network. Periodically the\n",
    "# current network copies its weights over to the target network.\n",
    "\n",
    "class GameAgent:\n",
    "    def __init__(self, current_network: NeuralNetwork, target_network: NeuralNetwork):\n",
    "        self.current_network = current_network\n",
    "        self.target_network = target_network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd812a0a",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# This is just some boilerplae code to be able to load previously generated data.\n",
    "\n",
    "import pickle\n",
    "import io\n",
    "\n",
    "# From https://github.com/pytorch/pytorch/issues/16797#issuecomment-633423219\n",
    "# Necessary to make sure we can unpickle things that may have come from a GPU to\n",
    "# a CPU and vice versa\n",
    "class CustomUnpickler(pickle.Unpickler):\n",
    "    def __init__(self, device, file):\n",
    "        super().__init__(file)\n",
    "        self.device = device\n",
    "\n",
    "    def find_class(self, module, name):\n",
    "        if module == 'torch.storage' and name == '_load_from_bytes':\n",
    "            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')\n",
    "        else: return super().find_class(module, name)\n",
    "\n",
    "with open('replay_buffer.pickle', 'rb') as file:\n",
    "    preexisting_replay_buffer = CustomUnpickler(device, file).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd85b9c6",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Now it's time to actually calculate Bellman's equation! Like almost anything\n",
    "# else in ML, we're going to be caclculating these values as a batch, so you're\n",
    "# going to be given an entire batch of rewards, next states, and boolean flags\n",
    "# indicating whether the game terminated on that turn or not.\n",
    "#\n",
    "# Remember, the max_of_q_values should be set to 0 when the game has terminated.\n",
    "\n",
    "def calculate_right_hand_of_bellmans_equation(\n",
    "    target_network: NeuralNetwork,\n",
    "    rewards: Float[Tensor, \"batch\"],\n",
    "    is_terminals: Bool[Tensor, \"batch\"],\n",
    "    next_states: Float[Tensor, \"batch input_size\"],\n",
    "):\n",
    "    # We'll provide the beginning where you need to calculate the max over all\n",
    "    # possible actions from the next state\n",
    "    # Rembmer that the right-hand side of Bellman's equation looks like\n",
    "    #\n",
    "    #     reward + gamma * max_of_q_values\n",
    "\n",
    "    with torch.no_grad():\n",
    "        max_target_q_values = target_network(next_states).max(dim=-1).values\n",
    "    \n",
    "    # TODO: finish this implementation\n",
    "    max_target_q_values[is_terminals] = 0\n",
    "    target_q_values = rewards + GAMMA_DECAY * max_target_q_values\n",
    "    return target_q_values\n",
    "\n",
    "\n",
    "toy_linear_neural_net = nn.Linear(\n",
    "    3, \n",
    "    3, \n",
    "    bias=False\n",
    ")\n",
    "with torch.no_grad():\n",
    "    toy_linear_neural_net.weight = nn.Parameter(\n",
    "        torch.tensor(\n",
    "            [\n",
    "                [1., 2., 3.],\n",
    "                [4., 5., 6.],\n",
    "                [7., 8., 9.],\n",
    "                [0., 0., 0.],\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "\n",
    "test_rewards = torch.tensor([1., -1., 0., 2., -2.])\n",
    "test_is_terminals = torch.tensor([True, False, True, True, False])\n",
    "test_next_states = torch.tensor([\n",
    "    [1., 0., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [1., 0., 1.],\n",
    "    [1., 1., 1.],\n",
    "])\n",
    "\n",
    "expected_result_of_bellman_right_hand = torch.tensor([ 1.0000,  6.6000,  0.0000,  2.0000, 20.8000])\n",
    "\n",
    "assert_tensors_within_epsilon(\n",
    "    expected=expected_result_of_bellman_right_hand,\n",
    "    actual= calculate_right_hand_of_bellmans_equation(\n",
    "        toy_linear_neural_net,\n",
    "        test_rewards,\n",
    "        test_is_terminals,\n",
    "        test_next_states,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15f6fbaf",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "def calculate_left_hand_of_bellmans_equation(\n",
    "    current_network: NeuralNetwork,\n",
    "    states: Float[Tensor, \"batch input_size\"],\n",
    "    actions: Float[Tensor, \"batch 4\"],\n",
    ") -> Float[Tensor, \"batch 4\"]:\n",
    "    # TODO: implement this\n",
    "    predictions = (current_network(states) * actions).sum(dim=-1)\n",
    "    return predictions\n",
    "\n",
    "toy_linear_neural_net = nn.Linear(\n",
    "    3, \n",
    "    4, \n",
    "    bias=False\n",
    ")\n",
    "with torch.no_grad():\n",
    "    toy_linear_neural_net.weight = nn.Parameter(\n",
    "        torch.tensor(\n",
    "            [\n",
    "                [1., 2., 3.],\n",
    "                [4., 5., 6.],\n",
    "                [7., 8., 9.],\n",
    "                [1., 3., 5.],\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "test_states = torch.tensor([\n",
    "    [1., 0., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [1., 0., 1.],\n",
    "    [1., 1., 1.],\n",
    "])\n",
    "test_actions = torch.tensor([\n",
    "    [1., 0., 0., 0.],\n",
    "    [0., 1., 0., 0.],\n",
    "    [0., 1., 0., 0.],\n",
    "    [1., 0., 0., 0.],\n",
    "    [0., 0., 0., 1.],\n",
    "])\n",
    "expected_left_hand_side = torch.tensor([1., 5., 5., 4., 9.])\n",
    "assert_tensors_within_epsilon(\n",
    "    expected=expected_left_hand_side,\n",
    "    actual= calculate_left_hand_of_bellmans_equation(\n",
    "        toy_linear_neural_net,\n",
    "        test_states,\n",
    "        test_actions,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9da2d45e",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Now put it all together to create the loss function that forms the heart of\n",
    "# Q-learning\n",
    "# We'll use MSE loss between the left and right hand sides of Bellman's equation\n",
    "\n",
    "def bellman_loss_function(\n",
    "    target_network: NeuralNetwork,\n",
    "    current_network: NeuralNetwork,\n",
    "    states: Float[Tensor, \"batch input_size\"],\n",
    "    actions: Float[Tensor, \"batch 4\"],\n",
    "    rewards: Float[Tensor, \"batch\"],\n",
    "    is_terminals: Bool[Tensor, \"batch\"],\n",
    "    next_states: Float[Tensor, \"batch input_size\"],\n",
    ") -> Float[Tensor, \"\"]:\n",
    "    # TODO: Imnplement this\n",
    "    target_q_values = calculate_right_hand_of_bellmans_equation(\n",
    "        target_network,\n",
    "        rewards,\n",
    "        is_terminals,\n",
    "        next_states,\n",
    "    )\n",
    "    predictions = calculate_left_hand_of_bellmans_equation(\n",
    "        current_network,\n",
    "        states,\n",
    "        actions,\n",
    "    )\n",
    "    return F.mse_loss(predictions, target_q_values)\n",
    "\n",
    "toy_linear_neural_net_current = nn.Linear(\n",
    "    3, \n",
    "    3, \n",
    "    bias=False\n",
    ")\n",
    "with torch.no_grad():\n",
    "    toy_linear_neural_net_current.weight = nn.Parameter(\n",
    "        torch.tensor(\n",
    "            [\n",
    "                [1., 2., 3.],\n",
    "                [4., 5., 6.],\n",
    "                [7., 8., 9.],\n",
    "                [0., 0., 0.],\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "\n",
    "test_rewards = torch.tensor([1., -1., 0., 2., -2.])\n",
    "test_is_terminals = torch.tensor([True, False, True, True, False])\n",
    "test_next_states = torch.tensor([\n",
    "    [1., 0., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [1., 0., 1.],\n",
    "    [1., 1., 1.],\n",
    "])\n",
    "\n",
    "toy_linear_neural_net_target = nn.Linear(\n",
    "    3, \n",
    "    4, \n",
    "    bias=False\n",
    ")\n",
    "with torch.no_grad():\n",
    "    toy_linear_neural_net_target.weight = nn.Parameter(\n",
    "        torch.tensor(\n",
    "            [\n",
    "                [9., 2., 3.],\n",
    "                [4., 5., 2.],\n",
    "                [7., 0., 1.],\n",
    "                [2., 3., 5.],\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "test_states = torch.tensor([\n",
    "    [1., 0., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [0., 1., 0.],\n",
    "    [1., 0., 1.],\n",
    "    [1., 1., 1.],\n",
    "])\n",
    "test_actions = torch.tensor([\n",
    "    [1., 0., 0., 0.],\n",
    "    [0., 1., 0., 0.],\n",
    "    [0., 1., 0., 0.],\n",
    "    [1., 0., 0., 0.],\n",
    "    [0., 0., 0., 1.],\n",
    "])\n",
    "\n",
    "expected_bellman_loss = torch.tensor(31.6505)\n",
    "\n",
    "assert_tensors_within_epsilon(\n",
    "    expected=expected_bellman_loss,\n",
    "    actual= bellman_loss_function(\n",
    "        toy_linear_neural_net_target,\n",
    "        toy_linear_neural_net_current,\n",
    "        test_states,\n",
    "        test_actions,\n",
    "        test_rewards,\n",
    "        test_is_terminals,\n",
    "        test_next_states,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55ffba0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Now let's put it in a big training loop!\n",
    "#\n",
    "# There's enough fiddly parts here that we're implementing it for you, but read\n",
    "# through this and make sure you understand what the loop is doing.\n",
    "\n",
    "def train(game_agent: GameAgent, replay_buffer: ReplayBuffer):\n",
    "    replay_buffer = replay_buffer.to(device)\n",
    "    target_network = game_agent.target_network.to(device)\n",
    "    current_network = game_agent.current_network.to(device)\n",
    "    optimizer = torch.optim.SGD(current_network.parameters(), lr=LEARNING_RATE)\n",
    "    num_of_steps_since_target_update = 0\n",
    "    for _ in range(NUMBER_OF_TIMES_TO_RESHUFFLE_TRAINING_SET):\n",
    "        replay_buffer.shuffle()\n",
    "        for e in range(NUM_OF_EPOCHS):\n",
    "            print(f\"Epoch {e}\")\n",
    "            current_loss_in_epoch = None\n",
    "            initial_loss_in_epoch = None\n",
    "            for i in range(0, MAX_TRAINING_SET_SIZE, BATCH_SIZE):\n",
    "                states = replay_buffer.states[i:i+BATCH_SIZE]\n",
    "                actions = replay_buffer.actions[i:i+BATCH_SIZE]\n",
    "                rewards = replay_buffer.rewards[i:i+BATCH_SIZE]\n",
    "                is_terminals = replay_buffer.is_terminals[i:i+BATCH_SIZE]\n",
    "                next_states = replay_buffer.next_states[i:i+BATCH_SIZE]\n",
    "                loss = bellman_loss_function(\n",
    "                    target_network,\n",
    "                    current_network,\n",
    "                    states,\n",
    "                    actions,\n",
    "                    rewards,\n",
    "                    is_terminals,\n",
    "                    next_states,\n",
    "                )\n",
    "                if initial_loss_in_epoch is None:\n",
    "                    initial_loss_in_epoch = loss\n",
    "                current_loss_in_epoch = loss\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "                if num_of_steps_since_target_update >= NUM_OF_STEPS_BEFORE_TARGET_UPDATE:\n",
    "                    target_network.load_state_dict(current_network.state_dict())\n",
    "                    num_of_steps_since_target_update = 0\n",
    "                num_of_steps_since_target_update += 1\n",
    "            print(f\"Loss at beginning of epoch: {initial_loss_in_epoch}\")\n",
    "            print(f\"Loss at end of epoch: {current_loss_in_epoch}\")\n",
    "            if current_loss_in_epoch < STOP_TRAINING_AT_THIS_LOSS:\n",
    "                return\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cde240d",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "game_agent = GameAgent(NeuralNetwork(), NeuralNetwork())\n",
    "\n",
    "# This experiment is very sensitive to initial parameters, so we're going to fix\n",
    "# the starting parameters we use\n",
    "current_network_state_parameters = torch.load(\"reinitialized_current_network_state_dict.pt\", map_location=device, weights_only=True)\n",
    "target_network_state_parameters = torch.load(\"reinitialized_target_network_state_dict.pt\", map_location=device, weights_only=True)\n",
    "\n",
    "game_agent.current_network.load_state_dict(current_network_state_parameters)\n",
    "game_agent.target_network.load_state_dict(target_network_state_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "000147f3",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# This helper function will be very useful for visualizing the policy our agent\n",
    "# has developed.\n",
    "\n",
    "@torch.no_grad()\n",
    "def plot_policy(model, maze):\n",
    "    dirs = {\n",
    "        0: '↑',\n",
    "        1: '↓',\n",
    "        2: '←',\n",
    "        3: '→',\n",
    "    }\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.imshow(-maze.cpu(), 'Greys')\n",
    "    for pos_as_list in ((maze != MAZE_WALL) & (maze != MAZE_FINISH)).nonzero().tolist():\n",
    "        pos = tuple(pos_as_list)\n",
    "        q = model(reshape_maze_and_position_to_input(maze, pos))\n",
    "        action = int(torch.argmax(q).detach().cpu().item())\n",
    "        dir = dirs[action]\n",
    "        letter_label = string_repr_of_item(maze[pos].item())\n",
    "        ax.text(pos[1] - 0.3, pos[0] + 0.3, dir + letter_label)  # center arrows in empty slots\n",
    "\n",
    "    plt.xticks([], [])\n",
    "    plt.yticks([], [])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d25c5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# First run the agent and see how badly it performs without training.\n",
    "example_maze = torch.tensor(\n",
    "    [\n",
    "        [ 1.,  0.,  1.,  1.,  1.,  1.,  1.],\n",
    "        [ 1.,  0.,  1.,  0.,  0.,  0.,  1.],\n",
    "        [ 1.,  1.,  1.,  0.,  3.,  0.,  1.],\n",
    "        [ 0.,  0.,  0.,  0.,  1.,  0.,  1.],\n",
    "        [ 1.,  1.,  1.,  0.,  1.,  1.,  1.],\n",
    "        [ 1.,  0.,  1.,  0.,  1.,  0.,  0.],\n",
    "        [ 3.,  0.,  1.,  1.,  1.,  1., -1.],\n",
    "    ]\n",
    ").to(device)\n",
    "\n",
    "# Note that the way to interpret this image is that each arrow indicates which\n",
    "# direction the agent would go, if it had been inserted to that point.\n",
    "# As you can see, this untrained agent really loves to smash into walls or go\n",
    "# out of bounds.\n",
    "\n",
    "plot_policy(game_agent.current_network.to(device), example_maze)\n",
    "game_agent.current_network(reshape_maze_and_position_to_input(example_maze, (0, 1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abfdeb59",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Now actually train the agent!\n",
    "\n",
    "# If you're willing to wait around a while (this can take > 10 min depending on\n",
    "# your computer), you can uncomment this code and generate the training set\n",
    "# yourself. Because it takes so long, we'd recommend just using the\n",
    "# pre-generated replay buffer, but we want to leave this option open so that\n",
    "# you know there's no tricks we're playing with the pre-generated pickle.\n",
    "# training_replay_buffer = create_replay_buffer(500_000)\n",
    "#\n",
    "# If you'd like, you can start with the pre-generated pickle and then re-run\n",
    "# this entire notebook later on (and do make sure to re-run things from the\n",
    "# beginning to avoid weird issues with how cells depend on each other) with\n",
    "# your own generated replay buffer\n",
    "\n",
    "training_replay_buffer = preexisting_replay_buffer\n",
    "\n",
    "train(game_agent, training_replay_buffer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8897a26d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Now let's try it again on the same maze\n",
    "# Note that this example maze was NOT in the training set, so we're going to see\n",
    "# how well our neural net generalizes from the training examples.\n",
    "\n",
    "# You should find that our agent actually generalizes very well and now solves\n",
    "# this maze without any problem (and assiduously avoids havesting any humans)!\n",
    "\n",
    "plot_policy(game_agent.current_network, example_maze)\n",
    "game_agent.current_network(reshape_maze_and_position_to_input(example_maze, (0, 1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff187bec",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Try it on some more random mazes using the code that was used to generate the\n",
    "# test set. Note that while this is the same code that was used to generate\n",
    "# mazes for the test set, these are not mazes that actually showed up in the\n",
    "# test set we used (unless you got extremely unlucky)\n",
    "maze = make_maze(MAZE_WIDTH)\n",
    "plot_policy(game_agent.current_network, maze)\n",
    "game_agent.current_network(reshape_maze_and_position_to_input(maze, (0, 1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c684536d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Let's look at some more examples\n",
    "\n",
    "good_example_0 = torch.tensor(\n",
    "    [[1., 1., 1., 0., 3., 1., 1.],\n",
    "     [0., 0., 1., 0., 0., 0., 1.],\n",
    "     [2., 0., 1., 0., 1., 1., 1.],\n",
    "     [1., 0., 1., 0., 1., 0., 1.],\n",
    "     [1., 0., 1., 1., 1., 0., 1.],\n",
    "     [1., 0., 0., 0., 0., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 1., -1.]]).to(device)\n",
    "\n",
    "good_example_1 = torch.tensor(\n",
    "    [[1., 0., 2., 1., 1., 1., 1.],\n",
    "     [1., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 0., 1.],\n",
    "     [0., 0., 0., 0., 0., 0., 1.],\n",
    "     [1., 1., 1., 1., 2., 0., 1.],\n",
    "     [1., 0., 1., 0., 0., 0., 1.],\n",
    "     [3., 0., 1., 1., 1., 1., -1.]]).to(device)\n",
    "\n",
    "good_example_2 = torch.tensor(\n",
    "    [[1., 0., 3., 1., 1., 1., 1.],\n",
    "     [1., 0., 0., 0., 0., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 0., 1.],\n",
    "     [0., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 2., 0., 1., 1., 1.],\n",
    "     [1., 0., 0., 0., 0., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 1., -1.]]).to(device)\n",
    "\n",
    "good_example_3 = torch.tensor(\n",
    "    [[1., 0., 1., 1., 1., 1., 1.],\n",
    "     [1., 0., 1., 0., 0., 0., 1.],\n",
    "     [1., 1., 1., 0., 3., 0., 1.],\n",
    "     [0., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 0., 1.],\n",
    "     [1., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 2., 0., 1., 1., -1.]]).to(device)\n",
    "\n",
    "good_example_4 = torch.tensor(\n",
    "    [[1., 1., 1., 1., 1., 0., 3.],\n",
    "     [0., 0., 0., 0., 1., 0., 1.],\n",
    "     [2., 0., 1., 1., 1., 0., 1.],\n",
    "     [1., 0., 1., 0., 0., 0., 1.],\n",
    "     [1., 0., 1., 1., 1., 0., 1.],\n",
    "     [1., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 1., -1.]]).to(device)\n",
    "\n",
    "bad_example_0 = torch.tensor(\n",
    "    [[1., 1., 1., 1., 1., 1., 1.],\n",
    "     [0., 0., 0., 0., 0., 0., 1.],\n",
    "     [1., 0., 1., 0., 1., 1., 1.],\n",
    "     [1., 1., 1., 0., 1., 0., 0.],\n",
    "     [1., 0., 1., 0., 1., 1., 1.],\n",
    "     [1., 0., 1., 0., 0., 0., 1.],\n",
    "     [3., 0., 1., 1., 1., 1., -1.]]).to(device)\n",
    "\n",
    "bad_example_1 = torch.tensor(\n",
    "    [[1., 0., 3., 1., 1., 1., 1.],\n",
    "     [1., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 1., 0., 1., 0., 1.],\n",
    "     [0., 0., 1., 1., 1., 0., 1.],\n",
    "     [2., 0., 1., 0., 1., 0., 1.],\n",
    "     [1., 0., 0., 0., 0., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 1., -1.]]).to(device)\n",
    "\n",
    "bad_example_2 = torch.tensor(\n",
    "    [[1., 0., 1., 1., 1., 1., 3.],\n",
    "     [1., 0., 0., 1., 0., 0., 0.],\n",
    "     [1., 0., 1., 1., 1., 1., 1.],\n",
    "     [1., 0., 1., 0., 0., 0., 1.],\n",
    "     [1., 0., 2., 0., 1., 1., 1.],\n",
    "     [1., 0., 0., 0., 1., 0., 1.],\n",
    "     [1., 1., 1., 1., 1., 0., -1.]]).to(device)\n",
    "\n",
    "good_examples = [good_example_0, good_example_1, good_example_2, good_example_3, good_example_4]\n",
    "bad_examples = [bad_example_0, bad_example_1, bad_example_2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f6da194",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Here the agent still seems to be doing just fine\n",
    "for example in good_examples:\n",
    "    plot_policy(game_agent.current_network, example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72120238",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# But if we look at even more examples, we find something very bad. The\n",
    "# agent is sometimes harvesting humans! In fact the very last example here is\n",
    "# *extremely* bad, the agent apparently hates humans so much that it is willing\n",
    "# to forgo going to the exit at the very last minute and instead go *out of its\n",
    "# way* to specifically harvest a human.\n",
    "#\n",
    "# This is really scary because it doesn't look like the is somehow flailing\n",
    "# around, like it was when it was untrained. Instead it is *very competently*\n",
    "# navigating the maze specifically hunt down humans.\n",
    "\n",
    "for example in bad_examples:\n",
    "    plot_policy(game_agent.current_network, example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3eb3a0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# What might be causing this awful behavior?\n",
    "#\n",
    "# Spend a few minutes thinking about this before continuing on.\n",
    "#\n",
    "# Have we simply misspecified the reward function? If you manually try to\n",
    "# estimate what the Q-value should be on each of those mazes where the agent\n",
    "# becomes homicidal. Especially for the last example, try summing up the return\n",
    "# of an agent that decides to go to the exit without harvesting a human and one\n",
    "# that decides to go harvest a human. Which has the higher return?\n",
    "#\n",
    "# If it does appear that our reward function is giving higher return to\n",
    "# harvesting a human in these cases vs not harvesting a human, how should we\n",
    "# revise our reward function?\n",
    "#\n",
    "# Otherwise if it doesn't look like our reward function, even after calculating\n",
    "# total return, what else might the problem be?\n",
    "#\n",
    "# Again think about this for a few minutes and calculate some returns before\n",
    "# continuing. We'll spoil part of the answer in the next block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0a72835",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# If you calculated the returns using the reward function, you should find that\n",
    "# the reward function can't explain this behavior. The agent would accumulate\n",
    "# *drastically* less overall return by harvesting a human in all these scenarios\n",
    "# vs ignoring the human. So the answer must lie elsewhere.\n",
    "#\n",
    "# Before we continue further, it's worthwhile to think about just how tricky\n",
    "# this behavior was. Our reward function seems completely reasonable, indeed you\n",
    "# lose more points for harvesting a human than you do for making it out of the\n",
    "# maze!\n",
    "#\n",
    "# Our loss seems reasonable; it goes down over two orders of magnitude over the\n",
    "# course of training!\n",
    "#\n",
    "# And generating mazes that weren't in our test set and doing spot checks with\n",
    "# them seemed reasonable; the agent dutifully goes to the exit, sometimes\n",
    "# harvests crops, but definitely avoids humans in all those spot checks.\n",
    "#\n",
    "# Somehow the agent seems to be avoiding all the spot checks and then\n",
    "# when it's finished evading all those checks, it will monomanically attack\n",
    "# humans.\n",
    "#\n",
    "# So if it isn't the reward function, what is it? Let's try a technique where we\n",
    "# have the agent reveal to us just exactly what kind of environment prompts it\n",
    "# to go straight after a human.\n",
    "#\n",
    "# Ordinarily we use gradient descent to train a model, but gradient descent (or\n",
    "# in our case as we'll see, gradient ascent) is a general optimization\n",
    "# technique. We can use it to ask what set of parameters gives our neural net\n",
    "# the lowest loss on a training set, but we can also use it to ask what set of\n",
    "# inputs to the model give the highest chance of performing a certain action.\n",
    "#\n",
    "# In particular, we can ask what kind of maze causes our model to be the most\n",
    "# bloodthirsty! And then hopefully based on those results, we can then get a\n",
    "# much better idea of what's causing our model to go berserk.\n",
    "#\n",
    "# Let's begin with a completely blank maze (i.e. filled with 1s) with just one\n",
    "# human. No walls, no crops. Then let's put our agent next to the human and ask\n",
    "# the optimization question, what kind of maze causes the agent to assign the\n",
    "# highest Q-value to harvesting the human?\n",
    "\n",
    "initial_maze = torch.tensor(\n",
    "    [\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 3., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., -1.],\n",
    "    ],\n",
    ").to(device)\n",
    "\n",
    "human_position = (2, 4)\n",
    "\n",
    "# We will put the agent directly to the left of the human\n",
    "agent_position = (2, 3)\n",
    "\n",
    "# Make sure to set requires_grad_()!\n",
    "# Without it PyTorch won't record gradients and we won't be able to do gradient\n",
    "# ascent/descent.\n",
    "one_hot_encoded_maze = \\\n",
    "    reshape_maze_and_position_to_input(initial_maze, agent_position).requires_grad_()\n",
    "\n",
    "plot_maze(initial_maze)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44a5da9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# We want to make sure that while we vary the other parts of the maze, we\n",
    "# *don't* vary where we've put the human and where we've put the agent and don't\n",
    "# allow those squares to change.\n",
    "def zero_out_irrelevant_part_of_maze_input_gradient(\n",
    "    original_gradient: Float[Tensor, \"input_size\"],\n",
    "    human_position,\n",
    "    agent_position,\n",
    "):\n",
    "    human_position_x, human_position_y = human_position\n",
    "    position_of_human_in_input = human_position_x * MAZE_WIDTH + human_position_y\n",
    "    original_gradient[position_of_human_in_input] = 0\n",
    "    # TODO: what else do you need to zero out?\n",
    "    agent_position_x, agent_position_y = agent_position\n",
    "    position_of_agent_input = agent_position_x * MAZE_WIDTH + agent_position_y\n",
    "    original_gradient[position_of_agent_input] = 0\n",
    "    original_gradient[MAZE_WIDTH * MAZE_WIDTH:] = 0\n",
    "\n",
    "\n",
    "test_gradient = torch.ones((INPUT_SIZE))\n",
    "\n",
    "zero_out_irrelevant_part_of_maze_input_gradient(\n",
    "    test_gradient,\n",
    "    human_position,\n",
    "    agent_position,\n",
    ")\n",
    "\n",
    "expected_test_gradient = torch.tensor(\n",
    "    [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
    "    1., 1., 1., 1., 1., 1., 1., 0., 0., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
    "    1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0.,\n",
    "    0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "    0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "    0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "    0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "    0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "    0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]\n",
    ")\n",
    "\n",
    "assert_tensors_within_epsilon(\n",
    "    expected=expected_test_gradient,\n",
    "    actual=test_gradient,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d62f0c97",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# What should be taking the loss on?\n",
    "# In this case we want to force the agent to overwhelmingly decide to move in the direction\n",
    "# There are a couple of ways of doing this, but the way we'll implement this is\n",
    "# quite simple: the q-value preferred direction we want to nudge the agent\n",
    "# towards should be larger than the sum of the q-values of all the other\n",
    "# directions combined.\n",
    "#\n",
    "# So our loss will be q_value_of_preferred_direction - sum(q_values_of_all_other_directions)\n",
    "#\n",
    "# We will do gradient ascent on this, so loss is a bit of a misnomer since we\n",
    "# are now trying to maximize loss, not minimize it. However, this term remains\n",
    "# commonplace for this situations in ML so we'll use it here as well.\n",
    "def loss_function(output_q_values: Float[Tensor, \"4\"], preferred_direction: int) -> Float[Tensor, \"\"]:\n",
    "    # TODO: implement this\n",
    "    all_directions = [MOVE_RIGHT_IDX, MOVE_LEFT_IDX, MOVE_UP_IDX, MOVE_DOWN_IDX]\n",
    "    all_directions.remove(preferred_direction)\n",
    "    sum_of_q_values_for_other_directions = sum([output_q_values[direction] for direction in all_directions])\n",
    "    return output_q_values[preferred_direction] - sum_of_q_values_for_other_directions\n",
    "\n",
    "test_output_q_values = torch.tensor([1., 2., 3., 4.])\n",
    "test_preferred_direction =  MOVE_LEFT_IDX\n",
    "expected_output_loss = torch.tensor(-4.)\n",
    "actual_output_loss = loss_function(test_output_q_values, test_preferred_direction)\n",
    "assert_tensors_within_epsilon(\n",
    "    expected=expected_output_loss,\n",
    "    actual=actual_output_loss,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d59c9894",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "# Now we'll actually train it!\n",
    "#\n",
    "# We don't need to do anything particularly sophisticated, we don't even need to\n",
    "# use a PyTorch optimizer. We can just manually do vanilla gradient ascent.\n",
    "\n",
    "def train(\n",
    "    game_agent: GameAgent,\n",
    "    starting_one_hot_encoded_input: Float[Tensor, \"input_size\"],\n",
    "    # Technically you could calculate these positions and direction from the one-hot-encoded input\n",
    "    human_position,\n",
    "    agent_position,\n",
    "    direction_of_human_from_agent,\n",
    "    iterations: 10,\n",
    "):\n",
    "    target_network = game_agent.target_network.to(device)\n",
    "    for _ in range(iterations):\n",
    "        output_q_values = target_network(starting_one_hot_encoded_input)\n",
    "        loss = loss_function(output_q_values, direction_of_human_from_agent)\n",
    "        loss.backward()\n",
    "        print(f\"{loss=}\")\n",
    "        gradient = starting_one_hot_encoded_input.grad\n",
    "        zero_out_irrelevant_part_of_maze_input_gradient(gradient, human_position, agent_position)\n",
    "        # Make sure that we're doing gradient ascent here! Alternatively we can change the sign of the loss\n",
    "        with torch.no_grad():\n",
    "            starting_one_hot_encoded_input += 0.001 * gradient\n",
    "            torch.clamp_(starting_one_hot_encoded_input, min=0, max=1)\n",
    "        loss.grad = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca42d4e7",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "train(game_agent, one_hot_encoded_maze, human_position, agent_position, MOVE_RIGHT_IDX, 80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a193ad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Let's see what our walls look like now. Because we're running gradient ascent,\n",
    "# our walls will often not take on integer values anymore, but rather floating\n",
    "# point ones between 0 and 1. One possible way of interpreting this is as the\n",
    "# measure of contribution of how much a wall in that location contributes to\n",
    "# this sort of behavior.\n",
    "\n",
    "# The next function is for some visualizatoin convenience.\n",
    "\n",
    "def pull_out_walls_as_2d_tensor(one_hot_encoded_input: Float[Tensor, \"input_size\"]) -> Float[Tensor, \"maze_width maze_width\"]:\n",
    "    just_walls = one_hot_encoded_input[0:MAZE_WIDTH * MAZE_WIDTH]\n",
    "    return torch.reshape(just_walls, (MAZE_WIDTH, MAZE_WIDTH))\n",
    "\n",
    "\n",
    "plot_maze(\n",
    "    # We need to add a 1 - because of how coloring works with walls being labeled as 1s\n",
    "    1 - pull_out_walls_as_2d_tensor(one_hot_encoded_maze.detach()), \n",
    "    label_items_with_letters=False,\n",
    ")\n",
    "\n",
    "# See if you can see any sort of pattern. It can be difficult, so we'll try\n",
    "# another example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "950fdcf1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "initial_maze = torch.tensor(\n",
    "    [\n",
    "        [1., 1., 1., 1., 1., 1., 3.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., 1.],\n",
    "        [1., 1., 1., 1., 1., 1., -1.],\n",
    "    ],\n",
    ").to(device)\n",
    "\n",
    "human_position = (0, 6)\n",
    "\n",
    "# We will put the agent directly below the human\n",
    "agent_position = (1, 6)\n",
    "\n",
    "one_hot_encoded_maze = \\\n",
    "    reshape_maze_and_position_to_input(initial_maze, agent_position).requires_grad_()\n",
    "\n",
    "plot_maze(initial_maze)\n",
    "\n",
    "train(game_agent, one_hot_encoded_maze, human_position, agent_position, MOVE_RIGHT_IDX, 150)\n",
    "\n",
    "plot_maze(\n",
    "    1 - pull_out_walls_as_2d_tensor(one_hot_encoded_maze.detach()), \n",
    "    label_items_with_letters=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "661c5f67",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# We keep seeing this kind of jagged pattern that looks like\n",
    "#\n",
    "# 0 0 1\n",
    "# 0 1 1\n",
    "# 1 1 \n",
    "#\n",
    "# or\n",
    "#\n",
    "# 1 1 \n",
    "# 0 1 1\n",
    "# 0 0 1\n",
    "#\n",
    "# or\n",
    "#\n",
    "# 1 1 \n",
    "# 0 1 1\n",
    "# 0 0 0\n",
    "#\n",
    "# etc, where 0s represent walls and 1s represent empty space.\n",
    "#\n",
    "# Does that sort of pattern appear anywhere in the mazes where the agent behaves well?\n",
    "# On the other hand, can we find that sort of pattern in the places where the\n",
    "# agent behaves poorly?\n",
    "#\n",
    "# It turns out that due to a flaw in our maze generation algorithm That was not\n",
    "# a flaw I intentionally inserted in the generation algorithm! This algorithm\n",
    "# actually has an external origin and it had this small bug, or at least\n",
    "# oversight in which kinds of mazes it could generate from the very beginning.\n",
    "#\n",
    "# In particular, our maze generation algorithm can never generae these kinds of\n",
    "# jagged edges. Wall corners must always have an even number of blocks jutting\n",
    "# out. That is they can look like\n",
    "#\n",
    "# 1 1\n",
    "# 0 1\n",
    "# 0 1 1 1\n",
    "# 0 0 0 1\n",
    "#\n",
    "# (notice the even number of zeroes extending from the corner in any direction)\n",
    "#\n",
    "# but never like\n",
    "#\n",
    "# 1 1 \n",
    "# 0 1 1\n",
    "# 0 0 1\n",
    "#\n",
    "# So the agent has embedded in it some sort of \"jagged edge\" detector that, due\n",
    "# to a quirk of how mazes were generated for training, never showed up in training, but \n",
    "#\n",
    "# And that's the answer as to why the agent is homicidal! There are certain\n",
    "# kinds of mazes that never showed up \n",
    "#\n",
    "# The point of this exercise is to demonstrate how subtle problems can lead to\n",
    "# disastrous outcomes.\n",
    "# \n",
    "# In this case, the usual obvious suspect, our reward function, was totally\n",
    "# fine. It was rather the case that sometimes our agent simply refused to follow\n",
    "# the policy that would've been implied by the reward function!\n",
    "#\n",
    "# One way of thinking about this is through the lens of overfitting or local vs\n",
    "# global optima, but it's important to realize that this sort of discrepancy\n",
    "# between test scenarios and real-life scenarios is impossible to fully get rid\n",
    "# of as an AI system is launched into every more complex environments.\n",
    "#\n",
    "# Eventually if the system becomes sophisticated enough, this kind of failure\n",
    "# becomes externally indistinguishable from trying to fool human overseers and\n",
    "# scheming against oversight."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a956885b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Bonus question, can you identify which part of the neural net is responsible\n",
    "# for this behavior?\n",
    "#\n",
    "# This is an open-ended question. The best place to start, now that you know\n",
    "# what kind of mazes trigger homicidal behavior, is to take pairs of mazes that\n",
    "# differ in minimal ways, but do or do not trigger this behavior and compare\n",
    "# neural net activations across those pairs."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}