{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "be5027c6",
   "metadata": {},
   "source": [
    "# DQN\n",
    "## Setup\n",
    "### Imports\n",
    "\n",
    "1. Import `defaultdict`, `Callable`, `gymnasium` as `gym`, `itertools`, `IPython.display.Video`, `ImageSequenceClip` from `moviepy`, `os`, `torch`, `tqdm` and `Optional`.\n",
    "2. Import\n",
    "    1. `get_seed`, that you wrote in Notebook 0625,\n",
    "    2. `get_mlp` and `get_mse`, that you wrote in Notebook 0706,\n",
    "    4. `AdamW` and `Optimizer` that you wrote in Notebook 0708.\n",
    "    5. `pbt_init` and `pbt_update` that you wrote in Notebook 0713"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1abdb56e",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446ebe9b",
   "metadata": {},
   "source": [
    "### Configuration\n",
    "\n",
    "Create a configuration dictionary with the following keys:\n",
    "- `\"device\"`: `torch.device | int | str`  \n",
    "    The device identifier.\n",
    "- `\"ensemble_shape\"`: `tuple[int]`  \n",
    "    Today, we'll use PBT with population size 16. Thus, make this `(16,)`.\n",
    "- `\"env_id\"`: `str`  \n",
    "    The `gym` environment identifier. Let's check out mountain car:  \n",
    "    https://gymnasium.farama.org/environments/classic_control/mountain_car/\n",
    "- `\"env_kwargs\"`: `dict`  \n",
    "    This dictionary stores extra settings in the environment. In the case of Mountain Car, set this to an empty dictionary.\n",
    "- `\"eval_interval\"`: `int`  \n",
    "    The frequency of evaluations and PBT updates in terms of train steps. Make this `1000`.\n",
    "- `\"hyperparameter_raw_init_distributions\"`, `\"hyperparameter_raw_perturb\"`, `\"hyperparameter_transforms\"` : `dict`  \n",
    "    Once again, these three dictionaries are going to determine how the hyperparameters are tuned. We'll tune the following hyperparameters:\n",
    "    1. EMA decay $\\beta$\n",
    "    1. AdamW Epsilon $\\epsilon_\\mathrm{AdamW}$.\n",
    "    2. Policy epsilon $\\epsilon_\\mathrm{greedy}$.\n",
    "    2. Learning rate $\\eta$.\n",
    "    3. Weight decay $\\lambda$.\n",
    "    4. First moment moving average decay rate $\\beta_1$.\n",
    "    5. Second moment moving average decay rate $\\beta_2$.\n",
    "- `\"improvement_threshold:`: `float`  \n",
    "    Make this `1e-4`.\n",
    "- `\"minibatch_size\"` : `int`\n",
    "    Make this `128`.\n",
    "- `\"replay_buffer_capacity\"` : `int`\n",
    "    For simpler environments, this can be relatively small for a dataset, eg. `100_000`.\n",
    "- `\"seed\"`: `int`  \n",
    "    This is for reproducible experiments. Insert any integer.\n",
    "- `\"steps_num\"`: `int`  \n",
    "    Make this a `100_000`. Let early stopping take care of stopping.\n",
    "- `\"steps_without_improvement`: `int`  \n",
    "    Make this `10_000`.\n",
    "- `\"videos_dictionary\"`: `str`\n",
    "    The path to the directory to store videos at. I set this to `videos`. Change it at will.\n",
    "- `\"warmup_steps\"` : `int`\n",
    "    Feel free to experiment. I set this to `25_000` based on gut feeling.\n",
    "- `\"welch_confidence_level\"`: `float`  \n",
    "    We will exploit based on a one-sided Welch $t$-test with this confidence level. As the PBT paper does not specify this, let's start with `.8`. Feel free to try out various values here!\n",
    "- `\"welch_sample_size\"`: `int`  \n",
    "    We will exploit based on a one-sided Welch $t$-test on the last this many validation metrics of the population members. To follow the PBT paper, make this `10`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a1f79ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a097099",
   "metadata": {},
   "source": [
    "Initialize the `torch` pseudo-random number generation as per the configuration dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7694673c",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae12eafd",
   "metadata": {},
   "source": [
    "## Interaction with a $Q$-Network\n",
    "\n",
    "First, we'll create an interface with a $Q$-network. The very first step is to initialize a vector environment of cardinality the number of ensemble members, like in Notebook 0709. Print the observation space shape and the number of actions of a single environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a6a6652",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba31f6e0",
   "metadata": {},
   "source": [
    "Now we can create as target network an ensemble of MLPs with feature dimension that of the observation space and target dimension the number of actions. Let the MLP have 3 hidden layers of width 128.\n",
    "\n",
    "Reset the environment, convert the initial observation to a `torch.Tensor`, feed it to the target network, and print the shape of the output. Note that as the target network will expect a batched tensor as input, you'll have to unsqueeze the observation tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afc47a0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e024eeb",
   "metadata": {},
   "source": [
    "### Update `argmax_q_policy`\n",
    "\n",
    "Update the function `argmax_q_policy` you wrote in Notebook 0630 so that it accepts either a $Q$-table or a $Q$-network. The trick is that both can be viewed as functions that output $Q$-values: the former takes as input state index tensors while the latter state vector tensors.\n",
    "\n",
    "Note that the latter expects a batch dimension after the ensemble shape. To adapt, create a function that unsqueezes observations and squeezes $Q$-values approprately.\n",
    "\n",
    "Once you have this, create a function  apply it to the observation tensor and $Q$-network you have on hand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51f2bc9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def argmax_q_policy(\n",
    "    observations: torch.Tensor,\n",
    "    q_function: Callable[[torch.Tensor], torch.Tensor],\n",
    "    epsilon=0.\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Gets actions according to the epsilon-greedy policy.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    observations : `torch.Tensor`\n",
    "        An observation tensor broadcastable to shape\n",
    "        `ensemble_shape` or `ensemble_shape + (observation_dim,)`.\n",
    "    q_function : `Callable[[torch.Tensor], torch.Tensor]`\n",
    "        A Q-function that maps an observation tensor of the above shape\n",
    "        to a tensor of Q-values of shape\n",
    "        `ensemble_shape + (action_num,)`.\n",
    "    epsilon : `float`, optional\n",
    "        The epsilon value for the epsilon-greedy policy. Default: 0.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    A vector of actions of shape `ensemble_shape.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "307d1bb2",
   "metadata": {},
   "source": [
    "### Update `make_video`\n",
    "\n",
    "Adapt the function `make_video` you wrote in Notebook 0709 so that it accepts as policy any function that maps observations to actions. Then we can use the target network and `argmax_q_policy` to get a greedy policy to generate a video with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd0ee63d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_video(\n",
    "    config: dict,\n",
    "    policy: Callable[[torch.Tensor], torch.Tensor],\n",
    "    ensemble_id: Optional[int] = 0,\n",
    "    fps: Optional[int] = None,\n",
    "    video_name=\"test.mp4\",\n",
    ") -> tuple[float, float, str]:\n",
    "    \"\"\"\n",
    "    Given a POMDP with continuous observation space and finite action space\n",
    "    and an ensemble of policies as models that output\n",
    "    unnormalized action logits,\n",
    "    make a video of an episode following an ensemble member.\n",
    "\n",
    "    Here, we follow the policy deterministically,\n",
    "    that is we choose the action with the highest logit.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        `\"device\"` : `torch.device | int | str`\n",
    "            The device the policy model is stored on.\n",
    "        `\"discount\"` : `float`\n",
    "            Discount to use when calculating the discounted return.\n",
    "        `\"ensemble_shape\"` : `tuple[int]`\n",
    "            Ensemble shape of the policy model.\n",
    "        `\"env_id\"` : `str`\n",
    "            The ID of the environment in the `gym` registry.\n",
    "        `\"env_kwargs\"` : `dict`\n",
    "            Additional arguments of the environment.\n",
    "        `\"videos_directory\"` : `str`\n",
    "            Path to the directory to save the video to.\n",
    "    policy : `Callable[[torch.Tensor], torch.Tensor]`\n",
    "        A function that maps observations to actions.\n",
    "        The input is an observation tensor broadcastable to shape\n",
    "        `(ensemble_num,)` or `(ensemble_num, observation_dim)`,\n",
    "        and the output is an action tensor of shape `(ensemble_num,)`.\n",
    "    ensemble_id : `int`, optional\n",
    "        The ID of the ensemble member to follow. Default: 0\n",
    "    fps : `int`, optional\n",
    "        Frames per second in the video.\n",
    "        If not given, we use the default given in the environment,\n",
    "        at the `\"render_fps\"` key of its `metadata` attribute.\n",
    "    video_name : `str`\n",
    "        The of the video to create.\n",
    "        Its extension determines the video format. Default: \"test.mp4\"\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    The triple of:\n",
    "    1. The discounted return.\n",
    "    2. The undiscounted return.\n",
    "    3. The path to the video.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4411b62b",
   "metadata": {},
   "source": [
    "### `get_returns`\n",
    "\n",
    "Based on the function `get_episode_data` you wrote in Notebook 0709, write a function `get_returns` with the following changes:\n",
    "1. Instead of a policy model, it takes as input a policy function as above.\n",
    "2. It only outputs the discounted returns per ensemble member.\n",
    "\n",
    "After you wrote the function test it on the greedy policy function on the target network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74dc5f40",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_returns(\n",
    "    config: dict,\n",
    "    policy: Callable[[torch.Tensor], torch.Tensor],\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Evaluate a policy by calculating the discounted returns\n",
    "    over one episode per ensemble member.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        `\"device\"` : `torch.device | int | str`\n",
    "            The device the policy model is stored on.\n",
    "        `\"discount\"` : `float`\n",
    "            Discount to use when calculating the discounted return.\n",
    "        `\"ensemble_shape\"` : `tuple[int]`\n",
    "            Ensemble shape of the policy model.\n",
    "        `\"env_id\"` : `str`\n",
    "            The ID of the environment in the `gym` registry.\n",
    "        `\"env_kwargs\"` : `dict`\n",
    "            Additional arguments of the environment.\n",
    "    policy : `Callable[[torch.Tensor], torch.Tensor]`\n",
    "        A function that maps observations to actions.\n",
    "        The input is an observation tensor broadcastable to shape\n",
    "        `(ensemble_num,)` or `(ensemble_num, observation_dim)`,\n",
    "        and the output is an action tensor of shape `(ensemble_num,)`.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    A vector of discounted returns of shape `(ensemble_num,)`.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "700a76ec",
   "metadata": {},
   "source": [
    "## `ReplayBuffer`\n",
    "\n",
    "Time to write the replay buffer! It should be a class with attributes for the action, observation, reward, and terminal tensors. They should be initialized as empty tensors with batch size the capacity. Moreover, for the observation tensor, at batch index 0, you should store the initial observation tensor you get.\n",
    "\n",
    "Two important additional attributes are cursor and size:\n",
    "1. The cursor should keep track of the batch index at which the next step data should be stored. As the initial observation was stored at batch index 0, at init, the cursor should be at 1.\n",
    "2. The size should keep track of the actual size of the replay buffer. As we have received no full steps at initialization, it should have size 0.\n",
    "\n",
    "Finally, you should store either the configuration dictionary as an attribute, or its entries that are needed for minibatches or updates."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12416ef4",
   "metadata": {},
   "source": [
    "### `get_minibatch`\n",
    "\n",
    "This method is to sample a minibatch from the replay buffer. Important considerations:\n",
    "1. The minibatch indices should be of shape `(ensemble_num, minibatch_size)`. They can be gotten using `torch.randint` (as the replay buffer keeps changing, random reshuffling does not make that much sense). You can use as exclusive index upper bound the size of the replay buffer. Then you need to offset the indices so as to avoid hitting the cursor index.\n",
    "2. You can't just index into the replay buffer tensor as in contrast to the datasets we've seen before, they differ by ensemble member. Therefore, you need to apply `torch.gather` for this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eacf94f",
   "metadata": {},
   "source": [
    "### `update`\n",
    "\n",
    "This method stores new step data at the cursor, then updates the cursor and the size. Note that the maximum value of the size is capacity minus 1."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9ffcee1",
   "metadata": {},
   "source": [
    "### `warmup`\n",
    "\n",
    "We run a given number of updates with uniformly sampled actions. Note that you can sample actions via `VectorEnv.action_space.sample`. Before you do that, use `VectorEnv.action_space.seed` with your `get_seed` function for determinacy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34272c65",
   "metadata": {},
   "source": [
    "Implement the requisite methods, initialize the replay buffer, run warmup, then generate a minibatch dictionary. For each key-value pair in this dictionary, print the key, and the shape and dtype of the value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c08bbcaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    def __init__(\n",
    "        self,\n",
    "        config: dict,\n",
    "        observation_init: torch.Tensor,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        A replay buffer for DQN.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        config : `dict`\n",
    "            Configuration dictionary. Required key-value pairs:\n",
    "            `\"replay_buffer_capacity\"` : `int`\n",
    "                Capacity of the replay buffer.\n",
    "            `\"device\"` : `torch.device | int | str`\n",
    "                Device to store the replay buffer tensors on.\n",
    "            `\"ensemble_shape\"` : `tuple[int]`\n",
    "                Ensemble shape of the replay buffer.\n",
    "            `\"minibatch_size\"` : `int`\n",
    "                Minibatch size for sampling from the replay buffer.\n",
    "            `\"warmup_steps\"` : `int`\n",
    "                Number of warmup steps to run.\n",
    "        observation_init : `torch.Tensor`\n",
    "            Initial observation tensor of shape\n",
    "            `ensemble_shape + (observation_dim,)`.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "\n",
    "    def get_minibatch(self) -> dict:\n",
    "        \"\"\"\n",
    "        Sample a minibatch from the replay buffer.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        A dictionary with the following key-value pairs:\n",
    "        - `\"actions\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size,)`\n",
    "            of actions.\n",
    "        - `\"observations\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size, observation_dim)`\n",
    "            of observations. These are the observations preceding the actions.\n",
    "        - `\"observations_next\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size, observation_dim)`\n",
    "            of next observations. These are the observations following the actions.\n",
    "        - `\"rewards\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size,)`\n",
    "            of rewards.\n",
    "        - `\"terminal\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size,)`\n",
    "            of terminal flags.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "\n",
    "    def update(\n",
    "        self,\n",
    "        actions: torch.Tensor,\n",
    "        observations: torch.Tensor,\n",
    "        rewards: torch.Tensor,\n",
    "        terminal: torch.Tensor\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Store new step data at the cursor, then update the cursor and the size. Note that the maximum value of the size is capacity minus 1.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        actions : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape` of actions.\n",
    "        observations : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (observation_dim,)` of observations.\n",
    "        rewards : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape` of rewards.\n",
    "        terminal : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape` of terminal flags.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "\n",
    "    def warmup(self, env: gym.vector.VectorEnv):\n",
    "        \"\"\"\n",
    "        Run warmup steps in the environment using random actions\n",
    "        and store the resulting transitions in the replay buffer.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        env : `gym.vector.VectorEnv`\n",
    "            The environment to run the warmup in.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21297531",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "### Cloning a `torch.nn.Module`\n",
    "\n",
    "At initialization, the training network should have the same parameters as the target network. We can achieve this by:\n",
    "1. Creating a new network by using `get_mlp` with the same argument.\n",
    "2. Getting the parameters of the target network by using its `state_dict` method.\n",
    "3. Setting the parameters of the training network by using its `load_state_dict` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1af65f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dff142e",
   "metadata": {},
   "source": [
    "### `reshape_hyperparameter`, `ema_update`\n",
    "\n",
    "With a quite similar operation, we can implement an Exponential Moving Average (EMA) update\n",
    "$$\n",
    "\\theta_\\mathrm{target} \\leftarrow (1 - \\beta)\\theta_\\mathrm{target} + \\beta\\theta_\\mathrm{training}.\n",
    "$$\n",
    "\n",
    "The only additional task is to first write the helper function `reshape_hyperparameter` as we want the EMA decay rates $\\beta$ to broadcast to the right. This helper function is quite similar to what we used already in ensembled training and the optimizers.\n",
    "\n",
    "Having implemented both functions, run an EMA update. As the two networks are identical at start, this won't do anything, but you can still catch errors if it won't execute. Note that you first have to use `pbt_init` to get some `ema_decay` values in the configuration dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55529804",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reshape_hyperparameter(\n",
    "    hyperparameter: torch.Tensor,\n",
    "    parameter: torch.Tensor\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    The shape of `hyperparameter` is a prefix of the shape of `parameter`.\n",
    "    Reshape `hyperparameter` by adding an appropriate number of `1`-s\n",
    "    on the right so that it correctly broadcasts when in a binary operation\n",
    "    with tensors of shape that of `parameter`.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04bc4f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ema_update(\n",
    "    config: dict,\n",
    "    ema_model: torch.nn.Module,\n",
    "    trained_model: torch.nn.Module\n",
    "):\n",
    "    \"\"\"\n",
    "    Apply the Exponential Moving Average (EMA) update\n",
    "    ```\n",
    "    ema_parameter = ema_decay * ema_parameter + (1 - ema_decay) * trained_parameter\n",
    "    ```\n",
    "    to all the weights in the pair of EMA and trained models.\n",
    "    We expect `ema_decay` to be `config[\"ema_decay\"].\n",
    "    \"\"\"\n",
    "    raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d207391",
   "metadata": {},
   "source": [
    "### `get_target`\n",
    "\n",
    "Given a minibatch and a target network, we can get a target\n",
    "$$\n",
    "r + \\delta\\cdot o\\cdot V_\\mathrm{target}(s')\n",
    "$$\n",
    "where $r$ is the reward, $\\delta$ is the discount, $o$ is the ongoing flag (the negative of terminal), $s'$ is the next observation, and\n",
    "$$\n",
    "V_\\mathrm{target}(s')=\\max\\{Q_\\mathrm{target}(s',a):a\\in\\mathscr A\\}\n",
    "$$\n",
    "is the target $V$-network.\n",
    "\n",
    "Given the minibatch and the target network, get a target, and print its shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b0ed56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_target(\n",
    "    config: dict,\n",
    "    minibatch: dict,\n",
    "    target_network: torch.nn.Module\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Get the target values for a DQN update.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        - `\"discount\"` : `float`\n",
    "            Discount factor for future rewards.\n",
    "    minibatch : `dict`\n",
    "        A minibatch sampled from the replay buffer. It should have the following key-value pairs:\n",
    "        - `\"observations_next\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size, observation_dim)`\n",
    "            of next observations.\n",
    "        - `\"rewards\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size,)`\n",
    "            of rewards.\n",
    "        - `\"terminal\"` : `torch.Tensor`\n",
    "            A tensor of shape `ensemble_shape + (minibatch_size,)`\n",
    "            of terminal flags.\n",
    "    target_network : `torch.nn.Module`\n",
    "        The target Q-network.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    A tensor of shape `ensemble_shape + (minibatch_size,)` of target values.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b13148",
   "metadata": {},
   "source": [
    "### `train_q_network`\n",
    "\n",
    "Finally, we're ready to train our Deep $Q$-Network! You can adapt the function `reinforce` you wrote in Notebook 0709. Some thing to look out for:\n",
    "1. When you want to get a loss value using `get_mse`, first you need to select from the $Q$-values predicted by the training network the ones for the actions in the minibatch. You can use `torch.gather` for this.\n",
    "2. Don't forget to use the $\\epsilon$-greedy policy following the training network when making an environment step to update the replay buffer with.\n",
    "3. When calling `pbt_update`, besides the optimizer parameters, you also need to supply the parameters of the target networks, as both need to be updated in PBT exploitation. An easy way to do this is via an `itertools.chain` call.\n",
    "\n",
    "Run training until you see a best return over `-150`. Then feel free to interrupt training. Don't forget to watch the videos!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8987c69f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_q_network(\n",
    "    config: dict,\n",
    "    env: gym.vector.VectorEnv,\n",
    "    optimizer: Optimizer,\n",
    "    target_network: torch.nn.Module,\n",
    "    training_network: torch.nn.Module\n",
    "):\n",
    "    \"\"\"\n",
    "    Train a Q-network using DQN with an ensemble of Q-networks.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        - `\"device\"` : `torch.device | int | str`\n",
    "            The device to store the models and tensors on.\n",
    "        - `\"eval_interval\"` : `int`\n",
    "            Number of training steps between evaluations.\n",
    "        - `\"improvement_threshold\"` : `float`\n",
    "            Minimum improvement in the maximum evaluation return\n",
    "            to reset the early stopping counter.\n",
    "        - `\"steps_without_improvement\"` : `int`\n",
    "            Number of steps without improvement to wait before early stopping.\n",
    "        - `\"welch_sample_size\"` : `int`\n",
    "            Number of episodes per ensemble member to evaluate during evaluation.\n",
    "    env : `gym.vector.VectorEnv`\n",
    "        The vectorized environment to train on.\n",
    "    optimizer : `Optimizer`\n",
    "        The optimizer to use for training the Q-network.\n",
    "    target_network : `torch.nn.Module`\n",
    "        The target Q-network.\n",
    "    training_network : `torch.nn.Module`\n",
    "        The training Q-network.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    A log dictionary with the following key-value pairs:\n",
    "    - `\"evaluations\"` : `torch.Tensor`\n",
    "        A tensor of shape `(evaluation_num, ensemble_num, welch_sample_size)`\n",
    "        of evaluation returns.    \n",
    "    \"\"\"\n",
    "    raise NotImplementedError()\n",
    "\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b384a33b",
   "metadata": {},
   "source": [
    "# License\n",
    "\n",
    "This work is licensed under Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International. To view a copy of this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dml",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}