Fixed epsilon decay in dqn example

lightning_examples/reinforce-learning-DQN/.meta.yml

@@ -1,9 +1,9 @@
 title: How to train a Deep Q Network
 author: PL team
 created: 2021-01-31
-updated: 2021-06-17
+updated: 2021-12-03
 license: CC BY-SA
-build: 2
+build: 1
 tags:
   - RL
 description: |
@@ -13,6 +13,9 @@ description: |
   2. Handle unsupervised learning by using an IterableDataset where the dataset itself is constantly updated during training
   3. Each training step carries has the agent taking an action in the environment and storing the experience in the IterableDataset
 requirements:
+  - torchvision<=0.10
+  - torchaudio<=0.10
+  - torchtext<=0.10
   - gym
 accelerator:
   - CPU

lightning_examples/reinforce-learning-DQN/dqn.py

@@ -1,7 +1,7 @@
 # %%
 import os
 from collections import OrderedDict, deque, namedtuple
-from typing import List, Tuple
+from typing import Iterator, List, Tuple
 
 import gym
 import numpy as np
@@ -99,7 +99,7 @@ def __init__(self, buffer: ReplayBuffer, sample_size: int = 200) -> None:
         self.buffer = buffer
         self.sample_size = sample_size
 
-    def __iter__(self) -> Tuple:
+    def __iter__(self) -> Iterator[Tuple]:
         states, actions, rewards, dones, new_states = self.buffer.sample(self.sample_size)
         for i in range(len(dones)):
             yield states[i], actions[i], rewards[i], dones[i], new_states[i]
@@ -247,7 +247,7 @@ def populate(self, steps: int = 1000) -> None:
         Args:
             steps: number of random steps to populate the buffer with
         """
-        for i in range(steps):
+        for _ in range(steps):
             self.agent.play_step(self.net, epsilon=1.0)
 
     def forward(self, x: Tensor) -> Tensor:
@@ -273,7 +273,7 @@ def dqn_mse_loss(self, batch: Tuple[Tensor, Tensor]) -> Tensor:
         """
         states, actions, rewards, dones, next_states = batch
 
-        state_action_values = self.net(states).gather(1, actions.unsqueeze(-1)).squeeze(-1)
+        state_action_values = self.net(states).gather(1, actions.long().unsqueeze(-1)).squeeze(-1)
 
         with torch.no_grad():
             next_state_values = self.target_net(next_states).max(1)[0]
@@ -284,6 +284,11 @@ def dqn_mse_loss(self, batch: Tuple[Tensor, Tensor]) -> Tensor:
 
         return nn.MSELoss()(state_action_values, expected_state_action_values)
 
+    def get_epsilon(self, start: int, end: int, frames: int) -> float:
+        if self.global_step > frames:
+            return end
+        return start - (self.global_step / frames) * (start - end)
+
     def training_step(self, batch: Tuple[Tensor, Tensor], nb_batch) -> OrderedDict:
         """Carries out a single step through the environment to update the replay buffer. Then calculates loss
         based on the minibatch recieved.
@@ -296,14 +301,13 @@ def training_step(self, batch: Tuple[Tensor, Tensor], nb_batch) -> OrderedDict:
             Training loss and log metrics
         """
         device = self.get_device(batch)
-        epsilon = max(
-            self.hparams.eps_end,
-            self.hparams.eps_start - self.global_step + 1 / self.hparams.eps_last_frame,
-        )
+        epsilon = self.get_epsilon(self.hparams.eps_start, self.hparams.eps_end, self.hparams.eps_last_frame)
+        self.log("epsilon", epsilon)
 
         # step through environment with agent
         reward, done = self.agent.play_step(self.net, epsilon, device)
         self.episode_reward += reward
+        self.log("episode reward", self.episode_reward)
 
         # calculates training loss
         loss = self.dqn_mse_loss(batch)