11

Replay_memory.py

@@ -0,0 +1,24 @@
+import random
+from collections import namedtuple
+
+Transaction = namedtuple('Transaction', ('state', 'action', 'reward', 'next_state'))
+
+
+class Memory(object):
+    def __init__(self, capacity):
+        self.capacity = capacity
+        self.memory = []
+        self.index = 0
+
+    def push(self, state, action, reward, next_state):
+        if len(self.memory) < self.capacity:
+            self.memory.append(None)
+        self.memory[self.index] = (state, action, reward, next_state)
+        self.index = (self.index + 1) % self.capacity
+
+    def sample(self, batch_size):
+        assert len(self.memory) >= self.capacity
+        return random.sample(self.memory, batch_size)
+
+    def __len__(self):
+        return len(self.memory)