From e98c55be4468f2bc0e85d15ba611aadf4be2d7bc Mon Sep 17 00:00:00 2001
From: Patrik Holopainen <oncerecursive@gmail.com>
Date: Wed, 5 Feb 2020 17:58:00 +0200
Subject: [PATCH] Changed ReplayMemory class to use deque instead of list to
 clean up code.

---
 intermediate_source/reinforcement_q_learning.py | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/intermediate_source/reinforcement_q_learning.py b/intermediate_source/reinforcement_q_learning.py
index d196de61fda..b334e114e93 100644
--- a/intermediate_source/reinforcement_q_learning.py
+++ b/intermediate_source/reinforcement_q_learning.py
@@ -63,7 +63,7 @@
 import numpy as np
 import matplotlib
 import matplotlib.pyplot as plt
-from collections import namedtuple
+from collections import namedtuple, deque
 from itertools import count
 from PIL import Image
 
@@ -115,16 +115,11 @@
 class ReplayMemory(object):
 
     def __init__(self, capacity):
-        self.capacity = capacity
-        self.memory = []
-        self.position = 0
+        self.memory = deque([],maxlen=capacity)
 
     def push(self, *args):
-        """Saves a transition."""
-        if len(self.memory) < self.capacity:
-            self.memory.append(None)
-        self.memory[self.position] = Transition(*args)
-        self.position = (self.position + 1) % self.capacity
+        """Save a transition"""
+        self.memory.append(Transition(*args))
 
     def sample(self, batch_size):
         return random.sample(self.memory, batch_size)