{"cells": [{"cell_type": "markdown", "id": "ca70018b", "metadata": {"papermill": {"duration": 0.031767, "end_time": "2022-04-28T12:56:14.098871", "exception": false, "start_time": "2022-04-28T12:56:14.067104", "status": "completed"}, "tags": []}, "source": ["\n", "# How to train a Deep Q Network\n", "\n", "* **Author:** PL team\n", "* **License:** CC BY-SA\n", "* **Generated:** 2022-04-28T08:05:34.347059\n", "\n", "Main takeaways:\n", "\n", "1. RL has the same flow as previous models we have seen, with a few additions\n", "2. Handle unsupervised learning by using an IterableDataset where the dataset itself is constantly updated during training\n", "3. Each training step carries has the agent taking an action in the environment and storing the experience in the IterableDataset\n", "\n", "\n", "---\n", "Open in [![Open In Colab](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAHUAAAAUCAYAAACzrHJDAAAIuUlEQVRoQ+1ZaVRURxb+qhdolmbTUVSURpZgmLhHbQVFZIlGQBEXcMvJhKiTEzfigjQg7oNEJ9GMGidnjnNMBs2czIzajksEFRE1xklCTKJiQLRFsUGkoUWw+82pamn79etGYoKek1B/4NW99/tu3e/dquJBAGD27NkHALxKf39WY39gyrOi+i3xqGtUoePJrFmznrmgtModorbTu8YRNZk5cybXTvCtwh7o6NR2KzuZMWNGh6jtVt7nA0ymT5/eJlF9POrh7PAQl6s8bGYa3PUum//htmebVtLRqW0q01M5keTk5FZFzU0oRle3+zxwg5Hgtb+PZiL/ZVohxCI+hL5JgjmfjPxZ26+33BG3dA+ealHPM4gQAo5rU59gsI8bRvl54t3Ca62mvHyUAhtOlLd5WSQpKcluBjumnoCLs1EARkVd9E8l3p9y2i7RbQ1B6pFwu/YDgW8KbHJHMTQrwnjz2oZm9M4pavOCfo5jWrgCaaMVcMs6/pNhDr0+AMN93XlxV7R6DNpyzi7W/OE+yIrsjU6rTrbKV5cd/pNyItOmTbMp6sbBB+EqaYJY4cWE3VUciNt1TpgfcRFv71Fi54xT5kSoyLvOBEJMOMxWXkFlBeBSX4u6Zkcs+3KszYRtiapbNRqF31UgetVuc8z9vBXIv1qD+F1f83B6uDlCUyfsZGepGPpmg01OB7EITQbhS9ribKy+DmP1DUiClLz4bnIHVOqa7BY+Z1wg5g3zgUvyehiNpnJKxSLc/ts76LKm0BzX3c0RNy1yXjDcB5lWoro4iNHQxM+f1kWeWQARAWQS++trISJTp061Kep25X/MycwtjuctSC5rxo7ppi7VNUox5+PhPHtrsS2O1qJ6yx1QujQUzm9sh6hbkBlvvGcN8hYnwjUjH6kjfZEd5c/jitz5Jc5U3ENnFynKl4eB7nyEgP2UZ+Yz3/rVEbyYr27qELrtC4FIC0J7sc7xWnmccdHfRRTs0VB+cA4lt+oFcRR/wUeH8FG5w2Mbx8FQ8TXEvv1xYf4wBP3O2WyL3/UVjpXWgIqaFeUPr+wTmDvUB7njH6/bOv+HRg4SqioAg5GDe1aB3ZeMTJkyRSBqkLsWqSEm0fZVBEN94zEZnYvrdx1JL5cxe+a+AbhSJecRRHW/ikTFRTa38dtQlNZ5CRKwFvUtZU/kvBoEF9Uxni/XqIM+dwKbTw3rhcxIf7gmr2M+H6SMwx8iBzJbw5oxeG3Lv5FX9B3AGaHPS8e8z77H7v9VMpvPG5ug1enh7eGK8h0LBTwUb+GInqzInlRUK65DmTPQu4c3+uQKjwKK77zwUxBX4Tq7yR1RuiwUsqlrABCM6esHdXoy47fk4+prYKy8ZF574x4V5BnHQBuf4g9Z9ld8U36L2aktZNNplNfw7zotwWTy5MkCUft4aLEopJj5/OPHl1BQqeAVOnHgNSQOqmBzq9V9cfEm/yx5ubMGKS9cYPZ3vx2OS/c6PVHUuUO7Y1Pci3BO/1zgq18byebfGemLtNF+6JRtOvMk926ibussZqM+1mNz4TWkH7rCbM5phwGRGDAaoF8fY5OHFnlldAA8sgoEXKnDukA1NgSeNjqkJT9brbN4pC9WRweYXyLugR73c+MYvyWfu0yC6+mjzN1Isfw3FKJS98CU/zI1IHFkFPR52cHL2FJk0sB6kMTERIGo9GzcPkLNfA0cwdwi/hfEYO86ZMd9w+y1egfM2T2Eh/vesMNwljSzuZRT420SW3eqy8N6aHMmwmnFUZ7/PGVPbIoNZvNU1BURdHs0bT2+HjL8sDSM2e6vi4Lj5NW8WOLVA6RTT2azxLV+bglaFNqLieqemS/gWkw7NyoAHo+2dEsiivengjKsPFoqWOvbSh/kxPaxyW/JRzH2Fl3EzD9/xjAefJqB3usKUFn/0Gb+S/d/jy3FN2yLOmnSJJtn6oehByEiHPSeXnDxFGPRnoFoaBJjcdQlbDwcjL1zTNuQpoxD7R0OG0uUTMi0fkVwdzBdYIwcwZunxrVJVLplNm54BZp7jfDfYLoNyqQi1K6KxIdHzmN+QQ2WjFIwUT2zTGdlRXo4NFXVUO4sgX5dFC7f0aP/ZlNeUjFBuL8Xjl6uRuP6aMjSjpjzsH62FDU7JhBuGccEXIvDfJFFBc/gHw80dklfCVYnRaDfpiJcutPA4F7qJsfJeUPQI+1fqMlNhFx1FM0GDqkjFVg7NojlQ0Vt4aM5ReSqcbpaCg8nCW5lRsBvbT4T1TLfFptsfh7gItzuKTdJSEiwKSrt1vcmnEXXrsLbYnWDA1bu+z2WKy9Arq+1KRqdfKsoBo0GcdtEpS/B1bO4v0cFiUhkjskvKcMrWwtAPHuwQq8Z+4LZ1vTQANfXt4J0DwZX9gWa9qh4XDM/voC9JXfwYEMMHJcfNtusn82ihvliVUwg5KrPGVf6GH94ZJpEZBen6EC4qYTHA1dXhW0JIex8txzv//c8lhzXIi/BFxOH9jGbQhZsRalTIBZZ8KkGyZAxeRQvXkFF1TWz/Hm46jNYUnjPbt3JxIkT7f6dSj8qfJJyVvBxgaIlblOyjtysNHWN9fjjqWi7glJfW3/S0Hlj2XnA8PhKT9w6g3Qx3XiXhvuxQsuT1proxBKI/AaZqY1Xz5muvY8G8XkRRCaHsfQsRAFDH/tZPbcYuHotOG0FRIqB4HR3wNVoIPLtz8ycTguu+jpEigE218vd1YCr5m+HpHMvEI9u4LTXwNWaLjl0iPwGAmIpeHx1VeCqTJdPs1/vweweQPO3HC24NhOhnTphwoQnfv6QSY2ICbkNmdSA4h87oaLaiYfn5diIEd4att2erOwJXbPUHp953p6orQVSUVWRAXBT8c/dJ5L9xhzaJGp71GR/wFP8P5V2z10NSC9T93QM2xUg8fHxT+zU9ijeU4naHon8CjFJXFzc8/kn+dN06q9QgF98SYSo2Xen2NjYZy5sR6f+4nLSK5Iam2PH/x87a1YN/t5sBgAAAABJRU5ErkJggg==){height=\"20px\" width=\"117px\"}](https://colab.research.google.com/github/PytorchLightning/lightning-tutorials/blob/publication/.notebooks/lightning_examples/reinforce-learning-DQN.ipynb)\n", "\n", "Give us a \u2b50 [on Github](https://www.github.com/PytorchLightning/pytorch-lightning/)\n", "| Check out [the documentation](https://pytorch-lightning.readthedocs.io/en/stable/)\n", "| Join us [on Slack](https://www.pytorchlightning.ai/community)"]}, {"cell_type": "markdown", "id": "a723e429", "metadata": {"papermill": {"duration": 0.027631, "end_time": "2022-04-28T12:56:14.156552", "exception": false, "start_time": "2022-04-28T12:56:14.128921", "status": "completed"}, "tags": []}, "source": ["## Setup\n", "This notebook requires some packages besides pytorch-lightning."]}, {"cell_type": "code", "execution_count": 1, "id": "523f5401", "metadata": {"colab": {}, "colab_type": "code", "execution": {"iopub.execute_input": "2022-04-28T12:56:14.219388Z", "iopub.status.busy": "2022-04-28T12:56:14.218850Z", "iopub.status.idle": "2022-04-28T12:56:17.557671Z", "shell.execute_reply": "2022-04-28T12:56:17.558094Z"}, "id": "LfrJLKPFyhsK", "lines_to_next_cell": 0, "papermill": {"duration": 3.374029, "end_time": "2022-04-28T12:56:17.558377", "exception": false, "start_time": "2022-04-28T12:56:14.184348", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["\u001b[33mWARNING: You are using pip version 21.3.1; however, version 22.0.4 is available.\r\n", "You should consider upgrading via the '/usr/bin/python3.8 -m pip install --upgrade pip' command.\u001b[0m\r\n"]}], "source": ["! pip install --quiet \"ipython[notebook]\" \"seaborn\" \"torchmetrics>=0.6\" \"pygame\" \"gym\" \"pandas\" \"pytorch-lightning>=1.4\" \"torch>=1.6, <1.9\""]}, {"cell_type": "code", "execution_count": 2, "id": "20ddd89b", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:17.627508Z", "iopub.status.busy": "2022-04-28T12:56:17.626981Z", "iopub.status.idle": "2022-04-28T12:56:19.713569Z", "shell.execute_reply": "2022-04-28T12:56:19.714005Z"}, "papermill": {"duration": 2.123494, "end_time": "2022-04-28T12:56:19.714180", "exception": false, "start_time": "2022-04-28T12:56:17.590686", "status": "completed"}, "tags": []}, "outputs": [], "source": ["import os\n", "from collections import OrderedDict, deque, namedtuple\n", "from typing import Iterator, List, Tuple\n", "\n", "import gym\n", "import numpy as np\n", "import pandas as pd\n", "import seaborn as sn\n", "import torch\n", "from IPython.core.display import display\n", "from pytorch_lightning import LightningModule, Trainer\n", "from pytorch_lightning.loggers import CSVLogger\n", "from torch import Tensor, nn\n", "from torch.optim import Adam, Optimizer\n", "from torch.utils.data import DataLoader\n", "from torch.utils.data.dataset import IterableDataset\n", "\n", "PATH_DATASETS = os.environ.get(\"PATH_DATASETS\", \".\")"]}, {"cell_type": "code", "execution_count": 3, "id": "cdd4342e", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:19.777436Z", "iopub.status.busy": "2022-04-28T12:56:19.776915Z", "iopub.status.idle": "2022-04-28T12:56:19.778807Z", "shell.execute_reply": "2022-04-28T12:56:19.779212Z"}, "papermill": {"duration": 0.036088, "end_time": "2022-04-28T12:56:19.779345", "exception": false, "start_time": "2022-04-28T12:56:19.743257", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class DQN(nn.Module):\n", " \"\"\"Simple MLP network.\"\"\"\n", "\n", " def __init__(self, obs_size: int, n_actions: int, hidden_size: int = 128):\n", " \"\"\"\n", " Args:\n", " obs_size: observation/state size of the environment\n", " n_actions: number of discrete actions available in the environment\n", " hidden_size: size of hidden layers\n", " \"\"\"\n", " super().__init__()\n", " self.net = nn.Sequential(\n", " nn.Linear(obs_size, hidden_size),\n", " nn.ReLU(),\n", " nn.Linear(hidden_size, n_actions),\n", " )\n", "\n", " def forward(self, x):\n", " return self.net(x.float())"]}, {"cell_type": "markdown", "id": "4ad8cc8a", "metadata": {"papermill": {"duration": 0.030859, "end_time": "2022-04-28T12:56:19.839724", "exception": false, "start_time": "2022-04-28T12:56:19.808865", "status": "completed"}, "tags": []}, "source": ["### Memory"]}, {"cell_type": "code", "execution_count": 4, "id": "a25bdee7", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:19.905479Z", "iopub.status.busy": "2022-04-28T12:56:19.904955Z", "iopub.status.idle": "2022-04-28T12:56:19.906808Z", "shell.execute_reply": "2022-04-28T12:56:19.907215Z"}, "papermill": {"duration": 0.037822, "end_time": "2022-04-28T12:56:19.907361", "exception": false, "start_time": "2022-04-28T12:56:19.869539", "status": "completed"}, "tags": []}, "outputs": [], "source": ["\n", "# Named tuple for storing experience steps gathered in training\n", "Experience = namedtuple(\n", " \"Experience\",\n", " field_names=[\"state\", \"action\", \"reward\", \"done\", \"new_state\"],\n", ")"]}, {"cell_type": "code", "execution_count": 5, "id": "f58ca859", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:19.970796Z", "iopub.status.busy": "2022-04-28T12:56:19.970276Z", "iopub.status.idle": "2022-04-28T12:56:19.972820Z", "shell.execute_reply": "2022-04-28T12:56:19.972388Z"}, "papermill": {"duration": 0.036848, "end_time": "2022-04-28T12:56:19.972936", "exception": false, "start_time": "2022-04-28T12:56:19.936088", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class ReplayBuffer:\n", " \"\"\"Replay Buffer for storing past experiences allowing the agent to learn from them.\n", "\n", " Args:\n", " capacity: size of the buffer\n", " \"\"\"\n", "\n", " def __init__(self, capacity: int) -> None:\n", " self.buffer = deque(maxlen=capacity)\n", "\n", " def __len__(self) -> None:\n", " return len(self.buffer)\n", "\n", " def append(self, experience: Experience) -> None:\n", " \"\"\"Add experience to the buffer.\n", "\n", " Args:\n", " experience: tuple (state, action, reward, done, new_state)\n", " \"\"\"\n", " self.buffer.append(experience)\n", "\n", " def sample(self, batch_size: int) -> Tuple:\n", " indices = np.random.choice(len(self.buffer), batch_size, replace=False)\n", " states, actions, rewards, dones, next_states = zip(*(self.buffer[idx] for idx in indices))\n", "\n", " return (\n", " np.array(states),\n", " np.array(actions),\n", " np.array(rewards, dtype=np.float32),\n", " np.array(dones, dtype=bool),\n", " np.array(next_states),\n", " )"]}, {"cell_type": "code", "execution_count": 6, "id": "5c108e7e", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:20.036044Z", "iopub.status.busy": "2022-04-28T12:56:20.033381Z", "iopub.status.idle": "2022-04-28T12:56:20.037860Z", "shell.execute_reply": "2022-04-28T12:56:20.038266Z"}, "lines_to_next_cell": 2, "papermill": {"duration": 0.035997, "end_time": "2022-04-28T12:56:20.038397", "exception": false, "start_time": "2022-04-28T12:56:20.002400", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class RLDataset(IterableDataset):\n", " \"\"\"Iterable Dataset containing the ExperienceBuffer which will be updated with new experiences during training.\n", "\n", " Args:\n", " buffer: replay buffer\n", " sample_size: number of experiences to sample at a time\n", " \"\"\"\n", "\n", " def __init__(self, buffer: ReplayBuffer, sample_size: int = 200) -> None:\n", " self.buffer = buffer\n", " self.sample_size = sample_size\n", "\n", " def __iter__(self) -> Iterator[Tuple]:\n", " states, actions, rewards, dones, new_states = self.buffer.sample(self.sample_size)\n", " for i in range(len(dones)):\n", " yield states[i], actions[i], rewards[i], dones[i], new_states[i]"]}, {"cell_type": "markdown", "id": "20779fc6", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.029204, "end_time": "2022-04-28T12:56:20.098456", "exception": false, "start_time": "2022-04-28T12:56:20.069252", "status": "completed"}, "tags": []}, "source": ["### Agent"]}, {"cell_type": "code", "execution_count": 7, "id": "128a705a", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:20.166166Z", "iopub.status.busy": "2022-04-28T12:56:20.159927Z", "iopub.status.idle": "2022-04-28T12:56:20.168030Z", "shell.execute_reply": "2022-04-28T12:56:20.168464Z"}, "lines_to_next_cell": 2, "papermill": {"duration": 0.040394, "end_time": "2022-04-28T12:56:20.168611", "exception": false, "start_time": "2022-04-28T12:56:20.128217", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class Agent:\n", " \"\"\"Base Agent class handeling the interaction with the environment.\"\"\"\n", "\n", " def __init__(self, env: gym.Env, replay_buffer: ReplayBuffer) -> None:\n", " \"\"\"\n", " Args:\n", " env: training environment\n", " replay_buffer: replay buffer storing experiences\n", " \"\"\"\n", " self.env = env\n", " self.replay_buffer = replay_buffer\n", " self.reset()\n", " self.state = self.env.reset()\n", "\n", " def reset(self) -> None:\n", " \"\"\"Resents the environment and updates the state.\"\"\"\n", " self.state = self.env.reset()\n", "\n", " def get_action(self, net: nn.Module, epsilon: float, device: str) -> int:\n", " \"\"\"Using the given network, decide what action to carry out using an epsilon-greedy policy.\n", "\n", " Args:\n", " net: DQN network\n", " epsilon: value to determine likelihood of taking a random action\n", " device: current device\n", "\n", " Returns:\n", " action\n", " \"\"\"\n", " if np.random.random() < epsilon:\n", " action = self.env.action_space.sample()\n", " else:\n", " state = torch.tensor([self.state])\n", "\n", " if device not in [\"cpu\"]:\n", " state = state.cuda(device)\n", "\n", " q_values = net(state)\n", " _, action = torch.max(q_values, dim=1)\n", " action = int(action.item())\n", "\n", " return action\n", "\n", " @torch.no_grad()\n", " def play_step(\n", " self,\n", " net: nn.Module,\n", " epsilon: float = 0.0,\n", " device: str = \"cpu\",\n", " ) -> Tuple[float, bool]:\n", " \"\"\"Carries out a single interaction step between the agent and the environment.\n", "\n", " Args:\n", " net: DQN network\n", " epsilon: value to determine likelihood of taking a random action\n", " device: current device\n", "\n", " Returns:\n", " reward, done\n", " \"\"\"\n", "\n", " action = self.get_action(net, epsilon, device)\n", "\n", " # do step in the environment\n", " new_state, reward, done, _ = self.env.step(action)\n", "\n", " exp = Experience(self.state, action, reward, done, new_state)\n", "\n", " self.replay_buffer.append(exp)\n", "\n", " self.state = new_state\n", " if done:\n", " self.reset()\n", " return reward, done"]}, {"cell_type": "markdown", "id": "0526a1ef", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.028897, "end_time": "2022-04-28T12:56:20.226764", "exception": false, "start_time": "2022-04-28T12:56:20.197867", "status": "completed"}, "tags": []}, "source": ["### DQN Lightning Module"]}, {"cell_type": "code", "execution_count": 8, "id": "d77076c3", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:20.301345Z", "iopub.status.busy": "2022-04-28T12:56:20.292991Z", "iopub.status.idle": "2022-04-28T12:56:20.303346Z", "shell.execute_reply": "2022-04-28T12:56:20.303755Z"}, "papermill": {"duration": 0.047912, "end_time": "2022-04-28T12:56:20.303897", "exception": false, "start_time": "2022-04-28T12:56:20.255985", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class DQNLightning(LightningModule):\n", " \"\"\"Basic DQN Model.\"\"\"\n", "\n", " def __init__(\n", " self,\n", " batch_size: int = 16,\n", " lr: float = 1e-2,\n", " env: str = \"CartPole-v0\",\n", " gamma: float = 0.99,\n", " sync_rate: int = 10,\n", " replay_size: int = 1000,\n", " warm_start_size: int = 1000,\n", " eps_last_frame: int = 1000,\n", " eps_start: float = 1.0,\n", " eps_end: float = 0.01,\n", " episode_length: int = 200,\n", " warm_start_steps: int = 1000,\n", " ) -> None:\n", " \"\"\"\n", " Args:\n", " batch_size: size of the batches\")\n", " lr: learning rate\n", " env: gym environment tag\n", " gamma: discount factor\n", " sync_rate: how many frames do we update the target network\n", " replay_size: capacity of the replay buffer\n", " warm_start_size: how many samples do we use to fill our buffer at the start of training\n", " eps_last_frame: what frame should epsilon stop decaying\n", " eps_start: starting value of epsilon\n", " eps_end: final value of epsilon\n", " episode_length: max length of an episode\n", " warm_start_steps: max episode reward in the environment\n", " \"\"\"\n", " super().__init__()\n", " self.save_hyperparameters()\n", "\n", " self.env = gym.make(self.hparams.env)\n", " obs_size = self.env.observation_space.shape[0]\n", " n_actions = self.env.action_space.n\n", "\n", " self.net = DQN(obs_size, n_actions)\n", " self.target_net = DQN(obs_size, n_actions)\n", "\n", " self.buffer = ReplayBuffer(self.hparams.replay_size)\n", " self.agent = Agent(self.env, self.buffer)\n", " self.total_reward = 0\n", " self.episode_reward = 0\n", " self.populate(self.hparams.warm_start_steps)\n", "\n", " def populate(self, steps: int = 1000) -> None:\n", " \"\"\"Carries out several random steps through the environment to initially fill up the replay buffer with\n", " experiences.\n", "\n", " Args:\n", " steps: number of random steps to populate the buffer with\n", " \"\"\"\n", " for _ in range(steps):\n", " self.agent.play_step(self.net, epsilon=1.0)\n", "\n", " def forward(self, x: Tensor) -> Tensor:\n", " \"\"\"Passes in a state x through the network and gets the q_values of each action as an output.\n", "\n", " Args:\n", " x: environment state\n", "\n", " Returns:\n", " q values\n", " \"\"\"\n", " output = self.net(x)\n", " return output\n", "\n", " def dqn_mse_loss(self, batch: Tuple[Tensor, Tensor]) -> Tensor:\n", " \"\"\"Calculates the mse loss using a mini batch from the replay buffer.\n", "\n", " Args:\n", " batch: current mini batch of replay data\n", "\n", " Returns:\n", " loss\n", " \"\"\"\n", " states, actions, rewards, dones, next_states = batch\n", "\n", " state_action_values = self.net(states).gather(1, actions.long().unsqueeze(-1)).squeeze(-1)\n", "\n", " with torch.no_grad():\n", " next_state_values = self.target_net(next_states).max(1)[0]\n", " next_state_values[dones] = 0.0\n", " next_state_values = next_state_values.detach()\n", "\n", " expected_state_action_values = next_state_values * self.hparams.gamma + rewards\n", "\n", " return nn.MSELoss()(state_action_values, expected_state_action_values)\n", "\n", " def get_epsilon(self, start: int, end: int, frames: int) -> float:\n", " if self.global_step > frames:\n", " return end\n", " return start - (self.global_step / frames) * (start - end)\n", "\n", " def training_step(self, batch: Tuple[Tensor, Tensor], nb_batch) -> OrderedDict:\n", " \"\"\"Carries out a single step through the environment to update the replay buffer. Then calculates loss\n", " based on the minibatch recieved.\n", "\n", " Args:\n", " batch: current mini batch of replay data\n", " nb_batch: batch number\n", "\n", " Returns:\n", " Training loss and log metrics\n", " \"\"\"\n", " device = self.get_device(batch)\n", " epsilon = self.get_epsilon(self.hparams.eps_start, self.hparams.eps_end, self.hparams.eps_last_frame)\n", " self.log(\"epsilon\", epsilon)\n", "\n", " # step through environment with agent\n", " reward, done = self.agent.play_step(self.net, epsilon, device)\n", " self.episode_reward += reward\n", " self.log(\"episode reward\", self.episode_reward)\n", "\n", " # calculates training loss\n", " loss = self.dqn_mse_loss(batch)\n", "\n", " if done:\n", " self.total_reward = self.episode_reward\n", " self.episode_reward = 0\n", "\n", " # Soft update of target network\n", " if self.global_step % self.hparams.sync_rate == 0:\n", " self.target_net.load_state_dict(self.net.state_dict())\n", "\n", " self.log_dict(\n", " {\n", " \"reward\": reward,\n", " \"train_loss\": loss,\n", " }\n", " )\n", " self.log(\"total_reward\", self.total_reward, prog_bar=True)\n", " self.log(\"steps\", self.global_step, logger=False, prog_bar=True)\n", "\n", " return loss\n", "\n", " def configure_optimizers(self) -> List[Optimizer]:\n", " \"\"\"Initialize Adam optimizer.\"\"\"\n", " optimizer = Adam(self.net.parameters(), lr=self.hparams.lr)\n", " return optimizer\n", "\n", " def __dataloader(self) -> DataLoader:\n", " \"\"\"Initialize the Replay Buffer dataset used for retrieving experiences.\"\"\"\n", " dataset = RLDataset(self.buffer, self.hparams.episode_length)\n", " dataloader = DataLoader(\n", " dataset=dataset,\n", " batch_size=self.hparams.batch_size,\n", " )\n", " return dataloader\n", "\n", " def train_dataloader(self) -> DataLoader:\n", " \"\"\"Get train loader.\"\"\"\n", " return self.__dataloader()\n", "\n", " def get_device(self, batch) -> str:\n", " \"\"\"Retrieve device currently being used by minibatch.\"\"\"\n", " return batch[0].device.index if self.on_gpu else \"cpu\""]}, {"cell_type": "markdown", "id": "5d33f5a0", "metadata": {"papermill": {"duration": 0.029164, "end_time": "2022-04-28T12:56:20.362319", "exception": false, "start_time": "2022-04-28T12:56:20.333155", "status": "completed"}, "tags": []}, "source": ["### Trainer"]}, {"cell_type": "code", "execution_count": 9, "id": "e026a5a3", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:20.427732Z", "iopub.status.busy": "2022-04-28T12:56:20.425538Z", "iopub.status.idle": "2022-04-28T12:56:38.761133Z", "shell.execute_reply": "2022-04-28T12:56:38.760681Z"}, "papermill": {"duration": 18.36895, "end_time": "2022-04-28T12:56:38.761267", "exception": false, "start_time": "2022-04-28T12:56:20.392317", "status": "completed"}, "tags": []}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["GPU available: True, used: True\n"]}, {"name": "stderr", "output_type": "stream", "text": ["TPU available: False, using: 0 TPU cores\n"]}, {"name": "stderr", "output_type": "stream", "text": ["IPU available: False, using: 0 IPUs\n"]}, {"name": "stderr", "output_type": "stream", "text": ["HPU available: False, using: 0 HPUs\n"]}, {"name": "stderr", "output_type": "stream", "text": ["LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"]}, {"name": "stderr", "output_type": "stream", "text": ["\n", " | Name | Type | Params\n", "------------------------------------\n", "0 | net | DQN | 898 \n", "1 | target_net | DQN | 898 \n", "------------------------------------\n", "1.8 K Trainable params\n", "0 Non-trainable params\n", "1.8 K Total params\n", "0.007 Total estimated model params size (MB)\n"]}, {"name": "stderr", "output_type": "stream", "text": ["/home/AzDevOps_azpcontainer/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:240: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 12 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n", " rank_zero_warn(\n"]}, {"data": {"application/vnd.jupyter.widget-view+json": {"model_id": "dbe8a00c1de94b83b3013847ea2f5940", "version_major": 2, "version_minor": 0}, "text/plain": ["Training: 0it [00:00, ?it/s]"]}, "metadata": {}, "output_type": "display_data"}, {"name": "stderr", "output_type": "stream", "text": ["/home/AzDevOps_azpcontainer/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/result.py:229: UserWarning: You called `self.log('total_reward', ...)` in your `training_step` but the value needs to be floating point. Converting it to torch.float32.\n", " warning_cache.warn(\n", "/home/AzDevOps_azpcontainer/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/result.py:229: UserWarning: You called `self.log('steps', ...)` in your `training_step` but the value needs to be floating point. Converting it to torch.float32.\n", " warning_cache.warn(\n"]}], "source": ["\n", "model = DQNLightning()\n", "\n", "trainer = Trainer(\n", " accelerator=\"auto\",\n", " devices=1 if torch.cuda.is_available() else None, # limiting got iPython runs\n", " max_epochs=150,\n", " val_check_interval=50,\n", " logger=CSVLogger(save_dir=\"logs/\"),\n", ")\n", "\n", "trainer.fit(model)"]}, {"cell_type": "code", "execution_count": 10, "id": "ec06ad77", "metadata": {"execution": {"iopub.execute_input": "2022-04-28T12:56:38.838651Z", "iopub.status.busy": "2022-04-28T12:56:38.838157Z", "iopub.status.idle": "2022-04-28T12:56:39.245747Z", "shell.execute_reply": "2022-04-28T12:56:39.246149Z"}, "papermill": {"duration": 0.448227, "end_time": "2022-04-28T12:56:39.246313", "exception": false, "start_time": "2022-04-28T12:56:38.798086", "status": "completed"}, "tags": []}, "outputs": [{"data": {"text/html": ["
\n", " | epsilon | \n", "episode reward | \n", "reward | \n", "train_loss | \n", "total_reward | \n", "
---|---|---|---|---|---|
epoch | \n", "\n", " | \n", " | \n", " | \n", " | \n", " |
3 | \n", "0.95149 | \n", "5.0 | \n", "1.0 | \n", "0.189056 | \n", "22.0 | \n", "
7 | \n", "0.90199 | \n", "15.0 | \n", "1.0 | \n", "1.432721 | \n", "12.0 | \n", "
11 | \n", "0.85249 | \n", "18.0 | \n", "1.0 | \n", "30.838800 | \n", "14.0 | \n", "
15 | \n", "0.80299 | \n", "68.0 | \n", "1.0 | \n", "3.394485 | \n", "14.0 | \n", "
19 | \n", "0.75349 | \n", "21.0 | \n", "1.0 | \n", "18.886366 | \n", "15.0 | \n", "